mnist_dpsgd_tutorial_keras_model.py

# Copyright 2021, The TensorFlow Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Training a CNN on MNIST with Keras and the DP SGD optimizer."""

fromabslimportapp
fromabslimportflags
fromabslimportlogging
importdp_accounting
importnumpyasnp
importtensorflowastf
fromtensorflow_privacy.privacy.keras_models.dp_keras_modelimportDPSequential


flags.DEFINE_boolean(
'dpsgd', True, 'If True, train with DP-SGD. If False, '
'train with vanilla SGD.')
flags.DEFINE_float('learning_rate', 0.15, 'Learning rate for training')
flags.DEFINE_float('noise_multiplier', 0.1,
'Ratio of the standard deviation to the clipping norm')
flags.DEFINE_float('l2_norm_clip', 1.0, 'Clipping norm')
flags.DEFINE_integer('batch_size', 250, 'Batch size')
flags.DEFINE_integer('epochs', 60, 'Number of epochs')
flags.DEFINE_integer(
'microbatches', 250, 'Number of microbatches '
'(must evenly divide batch_size)')
flags.DEFINE_string('model_dir', None, 'Model directory')

FLAGS=flags.FLAGS


defcompute_epsilon(steps):
"""Computes epsilon value for given hyperparameters."""
ifFLAGS.noise_multiplier==0.0:
returnfloat('inf')
orders= [1+x/10.forxinrange(1, 100)] +list(range(12, 64))
accountant=dp_accounting.rdp.RdpAccountant(orders)

sampling_probability=FLAGS.batch_size/60000
event=dp_accounting.SelfComposedDpEvent(
dp_accounting.PoissonSampledDpEvent(
sampling_probability,
dp_accounting.GaussianDpEvent(FLAGS.noise_multiplier)), steps)

accountant.compose(event)

# Delta is set to 1e-5 because MNIST has 60000 training points.
returnaccountant.get_epsilon(target_delta=1e-5)


defload_mnist():
"""Loads MNIST and preprocesses to combine training and validation data."""
train, test=tf.keras.datasets.mnist.load_data()
train_data, train_labels=train
test_data, test_labels=test

train_data=np.array(train_data, dtype=np.float32) /255
test_data=np.array(test_data, dtype=np.float32) /255

train_data=train_data.reshape((train_data.shape[0], 28, 28, 1))
test_data=test_data.reshape((test_data.shape[0], 28, 28, 1))

train_labels=np.array(train_labels, dtype=np.int32)
test_labels=np.array(test_labels, dtype=np.int32)

train_labels=tf.keras.utils.to_categorical(train_labels, num_classes=10)
test_labels=tf.keras.utils.to_categorical(test_labels, num_classes=10)

asserttrain_data.min() ==0.
asserttrain_data.max() ==1.
asserttest_data.min() ==0.
asserttest_data.max() ==1.

returntrain_data, train_labels, test_data, test_labels


defmain(unused_argv):
logging.set_verbosity(logging.INFO)
ifFLAGS.dpsgdandFLAGS.batch_size%FLAGS.microbatches!=0:
raiseValueError('Number of microbatches should divide evenly batch_size')

# Load training and test data.
train_data, train_labels, test_data, test_labels=load_mnist()

# Define a sequential Keras model
layers= [
tf.keras.layers.Conv2D(
16,
8,
strides=2,
padding='same',
activation='relu',
input_shape=(28, 28, 1)),
tf.keras.layers.MaxPool2D(2, 1),
tf.keras.layers.Conv2D(
32, 4, strides=2, padding='valid', activation='relu'),
tf.keras.layers.MaxPool2D(2, 1),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(32, activation='relu'),
tf.keras.layers.Dense(10)
 ]
ifFLAGS.dpsgd:
model=DPSequential(
l2_norm_clip=FLAGS.l2_norm_clip,
noise_multiplier=FLAGS.noise_multiplier,
num_microbatches=FLAGS.microbatches,
layers=layers,
 )
else:
model=tf.keras.Sequential(layers=layers)

optimizer=tf.keras.optimizers.SGD(learning_rate=FLAGS.learning_rate)
loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True)

# Compile model with Keras
model.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'])

# Train model with Keras
model.fit(
train_data,
train_labels,
epochs=FLAGS.epochs,
validation_data=(test_data, test_labels),
batch_size=FLAGS.batch_size)

# Compute the privacy budget expended.
ifFLAGS.dpsgd:
eps=compute_epsilon(FLAGS.epochs*60000//FLAGS.batch_size)
print('For delta=1e-5, the current epsilon is: %.2f'%eps)
else:
print('Trained with vanilla non-private SGD optimizer')


if__name__=='__main__':
app.run(main)