def forward_pass(self, batch, training=True):
del training # Unused
if isinstance(batch, dict):
batch = self.make_batch(**batch)
if not self._input_spec.y.is_compatible_with(batch.y):
raise ValueError('Expected batch.y to be compatible with '
'{} but found {}'.format(self._input_spec.y,
batch.y))
if not self._input_spec.x.is_compatible_with(batch.x):
raise ValueError('Expected batch.x to be compatible with '
'{} but found {}'.format(self._input_spec.x,
batch.x))
predictions = self._predict(batch.x)
residuals = predictions - batch.y
num_examples = tf.gather(tf.shape(predictions), 0)
total_loss = 0.5 * tf.reduce_sum(tf.pow(residuals, 2))
self._loss_sum.assign_add(total_loss)
self._num_examples.assign_add(num_examples)
self._num_batches.assign_add(1)
average_loss = total_loss / tf.cast(num_examples, tf.float32)
return model.BatchOutput(loss=average_loss,
predictions=predictions,
num_examples=num_examples)
def forward_pass(model_weights: ModelWeights,
batch_input: Any,
training: bool = True) -> model_lib.BatchOutput:
if isinstance(batch_input, collections.abc.Mapping):
x = batch_input['x']
y = batch_input['y']
elif isinstance(batch_input, collections.abc.Sequence):
x, y = batch_input
else:
raise ValueError(
'`batch_input` must be either a mapping with keys `x` '
f'and `y` or a sequence of `(x, y)`. Got: {batch_input!r}')
predictions = predict_on_batch(model_weights, x, training)
batch_loss = loss_fn(y_true=y, y_pred=predictions)
# TODO(b/207033265): more work needed to support models with multiple loss
# functions.
def nrows(t):
return t.nrows() if isinstance(t,
tf.RaggedTensor) else tf.shape(t)[0]
return model_lib.BatchOutput(loss=batch_loss,
predictions=predictions,
num_examples=nrows(
tf.nest.flatten(batch_input)[0]))
def _forward_pass(self, batch_input, training=True):
# forward_pass requires batch_input be a dictionary that can be passed to
# tf.keras.Model.__call__, namely it has keys `x`, and optionally `y`.
if hasattr(batch_input, '_asdict'):
batch_input = batch_input._asdict()
inputs = batch_input.get('x')
if inputs is None:
raise KeyError(
'Received a batch_input that is missing required key `x`. '
'Instead have keys {}'.format(list(batch_input.keys())))
predictions = self._keras_model(inputs=inputs, training=training)
y_true = batch_input.get('y')
if y_true is not None:
if len(self._loss_fns) == 1:
loss_fn = self._loss_fns[0]
batch_loss = loss_fn(y_true=y_true, y_pred=predictions)
else:
batch_loss = tf.zeros(())
for i in range(len(self._loss_fns)):
loss_fn = self._loss_fns[i]
loss_wt = self._loss_weights[i]
batch_loss += loss_wt * loss_fn(y_true=y_true[i],
y_pred=predictions[i])
else:
batch_loss = None
for metric in self.get_metrics():
metric.update_state(y_true=y_true, y_pred=predictions)
return model_lib.BatchOutput(loss=batch_loss, predictions=predictions)
def forward_pass(self, batch_input, training=True):
if training:
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True)
logits = self.predict_on_batch(batch_input['x'], training=True)
loss = loss_fn(y_true=batch_input['y'], y_pred=logits)
num_examples = tf.shape(logits)[0]
return model_lib.BatchOutput(loss=loss,
predictions=(),
num_examples=num_examples)
else:
predictions = self.predict_on_batch(batch_input['x'],
training=False)
return model_lib.BatchOutput(loss=(),
predictions=predictions,
num_examples=())
def forward_pass(self, batch, training=True):
assert not training
num_over = tf.reduce_sum(
tf.to_float(tf.greater(batch['temp'], self._variables.max_temp)))
tf.assign_add(self._variables.num_over, num_over)
loss = tf.constant(0.0)
predictions = tf.zeros_like(batch['temp'])
return model.BatchOutput(loss=loss, predictions=predictions)
def forward_pass(model_weights, batch_input, training):
"""Test forward_pass implementing linear regression on MSE."""
x, y = batch_input
predictions = predict_on_batch(model_weights, x, training)
residuals = predictions - y
num_examples = tf.shape(predictions)[0]
total_loss = tf.reduce_sum(tf.pow(residuals, 2.))
average_loss = total_loss / tf.cast(num_examples, tf.float32)
return model_lib.BatchOutput(
loss=average_loss, predictions=predictions, num_examples=num_examples)
def forward_pass(self, batch, training=True):
assert not training
num_over = tf.reduce_sum(
tf.cast(tf.greater(batch['temp'], self._variables.max_temp),
tf.float32))
self._variables.num_over.assign_add(num_over)
loss = tf.constant(0.0)
predictions = tf.zeros_like(batch['temp'])
return model.BatchOutput(loss=loss,
predictions=predictions,
num_examples=tf.shape(predictions)[0])
def _forward_pass(self, batch_input, training=True):
if hasattr(batch_input, '_asdict'):
batch_input = batch_input._asdict()
if isinstance(batch_input, collections.Mapping):
inputs = batch_input.get('x')
else:
inputs = batch_input[0]
if inputs is None:
raise KeyError(
'Received a batch_input that is missing required key `x`. '
'Instead have keys {}'.format(list(batch_input.keys())))
predictions = self._keras_model(inputs, training=training)
if isinstance(batch_input, collections.Mapping):
y_true = batch_input.get('y')
else:
y_true = batch_input[1]
if y_true is not None:
if len(self._loss_fns) == 1:
loss_fn = self._loss_fns[0]
# Note: we add each of the per-layer regularization losses to the loss
# that we use to update trainable parameters, in addition to the
# user-provided loss function. Keras does the same in the
# `tf.keras.Model` training step. This is expected to have no effect if
# no per-layer losses are added to the model.
batch_loss = tf.add_n(
[loss_fn(y_true=y_true, y_pred=predictions)] +
self._keras_model.losses)
else:
# Note: we add each of the per-layer regularization losses to the losses
# that we use to update trainable parameters, in addition to the
# user-provided loss functions. Keras does the same in the
# `tf.keras.Model` training step. This is expected to have no effect if
# no per-layer losses are added to the model.
batch_loss = tf.add_n([tf.zeros(())] +
self._keras_model.losses)
for i in range(len(self._loss_fns)):
loss_fn = self._loss_fns[i]
loss_wt = self._loss_weights[i]
batch_loss += loss_wt * loss_fn(y_true=y_true[i],
y_pred=predictions[i])
else:
batch_loss = None
# TODO(b/145308951): Follow up here to pass through sample_weight in the
# case that we have a model supporting masking.
for metric in self.get_metrics():
metric.update_state(y_true=y_true, y_pred=predictions)
return model_lib.BatchOutput(loss=batch_loss,
predictions=predictions,
num_examples=tf.shape(
tf.nest.flatten(inputs)[0])[0])
def _forward_pass(self, variables, batch_input, training=True):
if hasattr(batch_input, '_asdict'):
batch_input = batch_input._asdict()
if isinstance(batch_input, collections.abc.Mapping):
inputs = batch_input.get('x')
else:
inputs = batch_input[0]
if inputs is None:
raise KeyError(
'Received a batch_input that is missing required key `x`. '
'Instead have keys {}'.format(list(batch_input.keys())))
predictions = self._keras_model(inputs, training=training)
if isinstance(batch_input, collections.abc.Mapping):
y_true = batch_input.get('y')
else:
y_true = batch_input[1]
if y_true is not None:
if len(self._loss_fns) == 1:
loss_fn = self._loss_fns[0]
batch_loss = tf.add_n(
[loss_fn(y_true=y_true, y_pred=predictions)] +
self._keras_model.losses)
else:
batch_loss = tf.add_n([tf.zeros(())] +
self._keras_model.losses)
for i in range(len(self._loss_fns)):
loss_fn = self._loss_fns[i]
loss_wt = self._loss_weights[i]
batch_loss += loss_wt * loss_fn(y_true=y_true[i],
y_pred=predictions[i])
else:
batch_loss = None
# TODO(b/145308951): Follow up here to pass through sample_weight in the
# case that we have a model supporting masking.
num_examples = tf.cast(tf.size(y_true), tf.float32)
variables['num_examples'].assign_add(num_examples)
self._loss_metric.update_state(y_true=y_true, y_pred=predictions)
variables['loss'].assign_add(self._loss_metric.result() * num_examples)
for metric in self._metrics:
metric_value = metric.update_state(y_true=y_true,
y_pred=predictions)
variables[str(metric.name)].assign_add(metric.result() *
num_examples)
return model_lib.BatchOutput(loss=batch_loss,
predictions=predictions,
num_examples=tf.shape(
tf.nest.flatten(inputs)[0])[0])
def forward_pass(self, batch, training=True):
"""Unlike the TestModel implementation above, only tracks num_same."""
assert not training
# Calculate how many of the values in the training data match the input.
num_same = tf.reduce_sum(
tf.cast(tf.equal(batch['temp'], self._variables.given_nums),
tf.float32))
self._variables.num_same.assign_add(num_same)
# We're not actually training anything, so just use 0 loss and predictions.
loss = tf.constant(0.0)
predictions = tf.zeros_like(batch['temp'])
return model.BatchOutput(loss=loss,
predictions=predictions,
num_examples=tf.shape(predictions)[0])
def forward_pass(weights: functional.ModelWeights,
batch_input: Any,
training: bool = True) -> model_lib.BatchOutput:
x, y = batch_input
if not input_spec[1].is_compatible_with(y):
raise ValueError("Expected batch_input[1] to be compatible with "
f"{input_spec[1]} but found {y}")
if not input_spec[0].is_compatible_with(x):
raise ValueError("Expected batch_input[0] to be compatible with "
"{input_spec[0]} but found {x}")
predictions = predict_on_batch(weights, x=x, training=training)
residuals = predictions - y
num_examples = tf.gather(tf.shape(predictions), 0)
total_loss = tf.reduce_sum(tf.pow(residuals, 2))
average_loss = total_loss / tf.cast(num_examples, tf.float32)
return model_lib.BatchOutput(loss=average_loss,
predictions=predictions,
num_examples=num_examples)
def forward_pass(self, batch_input, training=True) -> model.BatchOutput:
if not self._input_spec['y'].is_compatible_with(batch_input['y']):
raise ValueError("Expected batch_input['y'] to be compatible with "
f"{self._input_spec['y']} but found {batch_input['y']}")
if not self._input_spec['x'].is_compatible_with(batch_input['x']):
raise ValueError("Expected batch_input['x'] to be compatible with "
"{self._input_spec['x']} but found {batch_input['x']}")
predictions = self.predict_on_batch(x=batch_input['x'], training=training)
residuals = predictions - batch_input['y']
num_examples = tf.gather(tf.shape(predictions), 0)
total_loss = 0.5 * tf.reduce_sum(tf.pow(residuals, 2))
self._loss_sum.assign_add(total_loss)
self._num_examples.assign_add(num_examples)
self._num_batches.assign_add(1)
average_loss = total_loss / tf.cast(num_examples, tf.float32)
return model.BatchOutput(
loss=average_loss, predictions=predictions, num_examples=num_examples)
def test_save_load_convert_to_tff_model(self, model_fn):
dataset = get_dataset()
functional_model = model_fn(input_spec=dataset.element_spec)
path = self.get_temp_dir()
serialization.save_functional_model(functional_model, path)
loaded_model = serialization.load_functional_model(path)
tff_model = functional.model_from_functional(loaded_model)
example_batch = next(iter(dataset))
for training in [True, False]:
self.assertAllClose(
tff_model.predict_on_batch(x=example_batch[0],
training=training), [[0.]] * 5)
for training in [True, False]:
tf.nest.map_structure(
lambda x, y: self.assertAllClose(x, y, atol=1e-2, rtol=1e-2),
tff_model.forward_pass(batch_input=example_batch,
training=training),
model_lib.BatchOutput(loss=74.250,
predictions=np.zeros(shape=[5, 1]),
num_examples=5))
def _forward_pass(self, batch_input, training=True):
if hasattr(batch_input, '_asdict'):
batch_input = batch_input._asdict()
if isinstance(batch_input, collections.Mapping):
inputs = batch_input.get('x')
else:
inputs = batch_input[0]
if inputs is None:
raise KeyError(
'Received a batch_input that is missing required key `x`. '
'Instead have keys {}'.format(list(batch_input.keys())))
predictions = self._keras_model(inputs=inputs, training=training)
if isinstance(batch_input, collections.Mapping):
y_true = batch_input.get('y')
else:
y_true = batch_input[1]
if y_true is not None:
if len(self._loss_fns) == 1:
loss_fn = self._loss_fns[0]
batch_loss = loss_fn(y_true=y_true, y_pred=predictions)
else:
batch_loss = tf.zeros(())
for i in range(len(self._loss_fns)):
loss_fn = self._loss_fns[i]
loss_wt = self._loss_weights[i]
batch_loss += loss_wt * loss_fn(y_true=y_true[i],
y_pred=predictions[i])
else:
batch_loss = None
for metric in self.get_metrics():
metric.update_state(y_true=y_true, y_pred=predictions)
return model_lib.BatchOutput(loss=batch_loss,
predictions=predictions,
num_examples=tf.shape(
tf.nest.flatten(inputs)[0])[0])
def forward_pass(self, batch_input, training=True):
return model_lib.BatchOutput(loss=0.0,
predictions=self.predict_on_batch,
num_examples=0)
