def test_keras_model_using_embeddings(self):
model = model_examples.build_embedding_keras_model()
dummy_batch = collections.OrderedDict(x=np.zeros([1]), y=np.zeros([1]))
tff_model = keras_utils.from_keras_model(
keras_model=model,
dummy_batch=dummy_batch,
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
metrics=[NumBatchesCounter(), NumExamplesCounter()])
# Create a batch with the size of the vocab. These examples will attempt to
# train the embedding so that the model produces
# i -> (i / output_size) + 5
input_vocab_size = 10
output_vocab_size = 5
xs = []
ys = []
for input_id in range(input_vocab_size):
xs.append(input_id)
ys.append((input_id / output_vocab_size + 5) % output_vocab_size)
batch = collections.OrderedDict(
x=np.expand_dims(np.array(xs, dtype=np.int64), axis=-1),
y=np.expand_dims(np.array(ys, dtype=np.int64), axis=-1))
num_train_steps = 3
for _ in range(num_train_steps):
batch_output = self.evaluate(tff_model.forward_pass(batch))
self.assertGreater(batch_output.loss, 0.0)
m = self.evaluate(tff_model.report_local_outputs())
self.assertEqual(m['num_batches'], [num_train_steps])
self.assertEqual(m['num_examples'], [input_vocab_size * num_train_steps])
self.assertGreater(m['loss'][0], 0.0)
self.assertEqual(m['loss'][1], input_vocab_size * num_train_steps)
def test_keras_model_using_embeddings(self):
model = model_examples.build_embedding_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=adam.Adam(),
loss=loss_fn,
metrics=[NumBatchesCounter(),
NumExamplesCounter()])
dummy_batch = collections.OrderedDict([
('x', np.zeros([1])),
('y', np.zeros([1])),
])
tff_model = model_utils.from_compiled_keras_model(
keras_model=model, dummy_batch=dummy_batch)
# Create a batch with the size of the vocab. These examples will attempt to
# train the embedding so that the model produces
# i -> (i / output_size) + 5
input_vocab_size = 10
output_vocab_size = 5
xs = []
ys = []
for input_id in range(input_vocab_size):
xs.append(input_id)
ys.append((input_id / output_vocab_size + 5) % output_vocab_size)
batch = {
'x': np.expand_dims(np.array(xs, dtype=np.int64), axis=-1),
'y': np.expand_dims(np.array(ys, dtype=np.int64), axis=-1),
}
prior_loss = float('inf')
num_iterations = 3
for _ in range(num_iterations):
r = self.evaluate(tff_model.train_on_batch(batch))
self.assertLess(r.loss, prior_loss)
prior_loss = r.loss
m = self.evaluate(tff_model.report_local_outputs())
self.assertEqual(m['num_batches'], [num_iterations])
self.assertEqual(m['num_examples'],
[input_vocab_size * num_iterations])
self.assertGreater(m['loss'][0], 0.0)
self.assertEqual(m['loss'][1], input_vocab_size * num_iterations)
