def test_keras_model_using_batch_norm(self): model = model_examples.build_conv_batch_norm_keras_model() def loss_fn(y_true, y_pred): loss_per_example = tf.keras.losses.sparse_categorical_crossentropy( y_true=y_true, y_pred=y_pred) return tf.reduce_mean(loss_per_example) model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=0.01), loss=loss_fn, metrics=[NumBatchesCounter(), NumExamplesCounter()]) dummy_batch = collections.OrderedDict([ ('x', np.zeros([1, 28 * 28], dtype=np.float32)), ('y', np.zeros([1, 1], dtype=np.int64)), ]) tff_model = keras_utils.from_compiled_keras_model( keras_model=model, dummy_batch=dummy_batch) batch_size = 2 batch = { 'x': np.random.uniform(low=0.0, high=1.0, size=[batch_size, 28 * 28]).astype(np.float32), 'y': np.random.random_integers(low=0, high=9, size=[batch_size, 1]).astype(np.int64), } num_iterations = 2 for _ in range(num_iterations): self.evaluate(tff_model.train_on_batch(batch)) m = self.evaluate(tff_model.report_local_outputs()) self.assertEqual(m['num_batches'], [num_iterations]) self.assertEqual(m['num_examples'], [batch_size * num_iterations]) self.assertGreater(m['loss'][0], 0.0) self.assertEqual(m['loss'][1], batch_size * num_iterations) # Ensure we can assign the FL trained model weights to a new model. tff_weights = model_utils.ModelWeights.from_model(tff_model) keras_model = model_examples.build_conv_batch_norm_keras_model() tff_weights.assign_weights_to(keras_model) def assert_all_weights_close(keras_weights, tff_weights): for keras_w, tff_w in zip(keras_weights, six.itervalues(tff_weights)): self.assertAllClose(self.evaluate(keras_w), self.evaluate(tff_w), atol=1e-4, msg='Variable [{}]'.format(keras_w.name)) assert_all_weights_close(keras_model.trainable_weights, tff_weights.trainable) assert_all_weights_close(keras_model.non_trainable_weights, tff_weights.non_trainable)
def test_keras_model_using_batch_norm(self): model = model_examples.build_conv_batch_norm_keras_model() input_spec = collections.OrderedDict(x=tf.TensorSpec( shape=[None, 28 * 28], dtype=tf.float32), y=tf.TensorSpec(shape=[None, 1], dtype=tf.int64)) tff_model = keras_utils.from_keras_model( keras_model=model, input_spec=input_spec, loss=tf.keras.losses.SparseCategoricalCrossentropy(), metrics=[NumBatchesCounter(), NumExamplesCounter()]) batch_size = 2 batch = collections.OrderedDict( x=np.random.uniform(low=0.0, high=1.0, size=[batch_size, 28 * 28]).astype(np.float32), y=np.random.random_integers(low=0, high=9, size=[batch_size, 1]).astype(np.int64)) num_train_steps = 2 for _ in range(num_train_steps): self.evaluate(tff_model.forward_pass(batch)) m = self.evaluate(tff_model.report_local_outputs()) self.assertEqual(m['num_batches'], [num_train_steps]) self.assertEqual(m['num_examples'], [batch_size * num_train_steps]) self.assertGreater(m['loss'][0], 0.0) self.assertEqual(m['loss'][1], batch_size * num_train_steps) # Ensure we can assign the FL trained model weights to a new model. tff_weights = model_utils.ModelWeights.from_model(tff_model) keras_model = model_examples.build_conv_batch_norm_keras_model() tff_weights.assign_weights_to(keras_model) def assert_all_weights_close(keras_weights, tff_weights): for keras_w, tff_w in zip(keras_weights, tff_weights): self.assertAllClose(self.evaluate(keras_w), self.evaluate(tff_w), atol=1e-4, msg='Variable [{}]'.format(keras_w.name)) assert_all_weights_close(keras_model.trainable_weights, tff_weights.trainable) assert_all_weights_close(keras_model.non_trainable_weights, tff_weights.non_trainable)