def test_dummy_batch_types(self, dummy_batch):
keras_model = model_examples.build_linear_regresion_keras_functional_model(
feature_dims=1)
tff_model = model_utils.from_keras_model(
keras_model=keras_model,
dummy_batch=dummy_batch,
loss=tf.keras.losses.MeanSquaredError())
self.assertIsInstance(tff_model, model_utils.EnhancedModel)
def _make_keras_model():
keras_model = model_examples.build_linear_regresion_keras_functional_model(
feature_dims)
keras_model.compile(
optimizer=gradient_descent.SGD(learning_rate=0.01),
loss=tf.keras.losses.MeanSquaredError(),
metrics=[NumBatchesCounter(),
NumExamplesCounter()])
return keras_model
def _train_loop():
keras_model = model_examples.build_linear_regresion_keras_functional_model(
feature_dims)
# If the model is intended to be used for training, it must be compiled.
keras_model.compile(
optimizer=gradient_descent.SGD(learning_rate=0.01),
loss=tf.keras.losses.MeanSquaredError(),
metrics=[NumBatchesCounter(),
NumExamplesCounter()])
tff_model = model_utils.from_compiled_keras_model(
keras_model=keras_model,
dummy_batch=_create_dummy_batch(feature_dims))
batch = {
'x':
np.array([[0.0] * feature_dims, [5.0] * feature_dims],
dtype=np.float32),
'y':
np.array([[0.0], [5.0 * feature_dims]], dtype=np.float32),
}
batch_output = tff_model.train_on_batch(batch)
with tf.control_dependencies(list(batch_output)):
metrics = tff_model.report_local_outputs()
return batch_output, metrics
