mnist_dpsgd_tutorial_eager.py

# Copyright 2019, 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 in TF Eager mode with DP-SGD optimizer."""

fromabslimportapp
fromabslimportflags
importdp_accounting
importnumpyasnp
importtensorflowastf
fromtensorflow_privacy.privacy.optimizers.dp_optimizerimportDPGradientDescentGaussianOptimizer

GradientDescentOptimizer=tf.compat.v1.train.GradientDescentOptimizer
tf.compat.v1.enable_eager_execution()

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', 1.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=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)


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

# Fetch the mnist data
train, test=tf.keras.datasets.mnist.load_data()
train_images, train_labels=train
test_images, test_labels=test

# Create a dataset object and batch for the training data
dataset=tf.data.Dataset.from_tensor_slices(
 (tf.cast(train_images[..., tf.newaxis] /255,
tf.float32), tf.cast(train_labels, tf.int64)))
dataset=dataset.shuffle(1000).batch(FLAGS.batch_size)

# Create a dataset object and batch for the test data
eval_dataset=tf.data.Dataset.from_tensor_slices(
 (tf.cast(test_images[..., tf.newaxis] /255,
tf.float32), tf.cast(test_labels, tf.int64)))
eval_dataset=eval_dataset.batch(10000)

# Define the model using tf.keras.layers
mnist_model=tf.keras.Sequential([
tf.keras.layers.Conv2D(
16, 8, strides=2, padding='same', activation='relu'),
tf.keras.layers.MaxPool2D(2, 1),
tf.keras.layers.Conv2D(32, 4, strides=2, activation='relu'),
tf.keras.layers.MaxPool2D(2, 1),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(32, activation='relu'),
tf.keras.layers.Dense(10)
 ])

# Instantiate the optimizer
ifFLAGS.dpsgd:
opt=DPGradientDescentGaussianOptimizer(
l2_norm_clip=FLAGS.l2_norm_clip,
noise_multiplier=FLAGS.noise_multiplier,
num_microbatches=FLAGS.microbatches,
learning_rate=FLAGS.learning_rate)
else:
opt=GradientDescentOptimizer(learning_rate=FLAGS.learning_rate)

# Training loop.
steps_per_epoch=60000//FLAGS.batch_size
forepochinrange(FLAGS.epochs):
# Train the model for one epoch.
for (_, (images, labels)) inenumerate(dataset.take(-1)):
withtf.GradientTape(persistent=True) asgradient_tape:
# This dummy call is needed to obtain the var list.
logits=mnist_model(images, training=True)
var_list=mnist_model.trainable_variables

# In Eager mode, the optimizer takes a function that returns the loss.
defloss_fn():
logits=mnist_model(images, training=True) # pylint: disable=undefined-loop-variable,cell-var-from-loop
loss=tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=labels, logits=logits) # pylint: disable=undefined-loop-variable,cell-var-from-loop
# If training without privacy, the loss is a scalar not a vector.
ifnotFLAGS.dpsgd:
loss=tf.reduce_mean(input_tensor=loss)
returnloss

ifFLAGS.dpsgd:
grads_and_vars=opt.compute_gradients(
loss_fn, var_list, gradient_tape=gradient_tape)
else:
grads_and_vars=opt.compute_gradients(loss_fn, var_list)

opt.apply_gradients(grads_and_vars)

# Evaluate the model and print results
for (_, (images, labels)) inenumerate(eval_dataset.take(-1)):
logits=mnist_model(images, training=False)
correct_preds=tf.equal(tf.argmax(input=logits, axis=1), labels)
test_accuracy=np.mean(correct_preds.numpy())
print('Test accuracy after epoch %d is: %.3f'% (epoch, test_accuracy))

# Compute the privacy budget expended so far.
ifFLAGS.dpsgd:
eps=compute_epsilon((epoch+1) *steps_per_epoch)
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)