update dp keras model

This commit is contained in:
pranav subramani 2021-01-08 00:22:44 -07:00
parent 9d871b28c1
commit 6982e027b5

View file

@ -27,7 +27,7 @@ def make_dp_model_class(cls):
noise = tf.random.normal(
tf.shape(input=summed_grads), stddev=noise_stddev)
noised_grads = summed_grads + noise
return noised_grads / tf.cast(stacked_grads.shape[0], noised_grads.dtype)
return noised_grads / tf.cast(stacked_grads.shape[0], tf.float32)
def compute_per_example_grads(self, data):
x, y = data
@ -47,7 +47,8 @@ def make_dp_model_class(cls):
y_pred, per_eg_loss, per_eg_grads = tf.vectorized_map(
self.compute_per_example_grads, data)
loss = tf.reduce_mean(per_eg_loss, axis=0)
grads = tf.nest.map_structure(self.reduce_per_example_grads, per_eg_grads)
grads = tf.nest.map_structure(
self.reduce_per_example_grads, per_eg_grads)
self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
self.compiled_metrics.update_state(y, y_pred)
return {m.name: m.result() for m in self.metrics}