def step_fn(features, labels):
with GradientTape() as tape:
predictions = model(features, training=True)
prediction_loss = keras.losses.mean_squared_error(
labels, predictions)
loss = math_ops.reduce_mean(prediction_loss)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads,
model.trainable_variables))
return loss
def _process_single_batch(model,
inputs,
targets,
output_loss_metrics=None,
sample_weights=None,
training=False):
"""Calculate the loss and gradient for one input batch.
The model weights are updated if training is set to True.
Arguments:
model: Model whose loss has to be calculated.
inputs: List of input arrays.
targets: List of target arrays.
output_loss_metrics: List of metrics that are used to aggregated output
loss values.
sample_weights: Optional list of sample weight arrays.
training: The boolean represents if the weights of the model are updated.
'fit' methods will set this to True while 'evaluate' methods will
set this to False.
Returns:
output of the model, total loss, the loss and the mask
associated with each output.
Raises:
ValueError: If the model has no loss to optimize.
"""
with backend.learning_phase_scope(1 if training else 0):
with GradientTape() as tape:
outs, loss, loss_metrics, aggregated_loss_metrics, masks\
= _model_loss(
model,
inputs,
targets,
output_loss_metrics=output_loss_metrics,
sample_weights=sample_weights,
training=training)
if loss is None:
raise ValueError('The model cannot be run '
'because it has no loss to optimize.')
if training:
if not model._collected_trainable_weights:
logging.warning(
'The list of trainable weights is empty. Make sure that'
' you are not setting model.trainable to False before '
'compiling the model.')
else:
grads = tape.gradient(loss, model._collected_trainable_weights)
model.optimizer.apply_gradients(
zip(grads, model._collected_trainable_weights))
return outs, loss, loss_metrics, aggregated_loss_metrics, masks
def _process_single_batch(eager_model_inputs,
eager_model_outputs,
model,
training=False):
"""Calculate the loss and gradient for one input batch.
The model weights are updated if training is set to True.
Arguments:
eager_model_inputs: Input batch data.
eager_model_outputs: Output batch data.
model: Model whose loss has to be calculated.
training: The boolean represents if the weights of the model are updated.
'fit' methods will set this to True while 'evaluate' methods will
set this to False.
Returns:
output of the model, total loss and the loss associated with each output.
Raises:
ValueError: If the model loss is 0 or if the trainable weights list is
empty when the trainable parameter is set to True.
"""
K.set_learning_phase(training)
with GradientTape() as tape:
outs, loss, loss_metrics = _model_loss(model,
eager_model_inputs,
eager_model_outputs,
training=training)
if loss is None:
raise ValueError('The model cannot be run '
'because it has no loss to optimize.')
if training:
if not model._collected_trainable_weights:
raise ValueError(
'The list of trainable weights is empty. Make sure that '
'you are not setting model.trainable to False before '
'compiling the model.')
grads = tape.gradient(loss, model._collected_trainable_weights)
model.optimizer.apply_gradients(
zip(grads, model._collected_trainable_weights))
return outs, loss, loss_metrics
def _process_single_batch(model,
inputs,
targets,
output_loss_metrics=None,
sample_weights=None,
training=False):
"""Calculate the loss and gradient for one input batch.
The model weights are updated if training is set to True.
Args:
model: Model whose loss has to be calculated.
inputs: List of input arrays.
targets: List of target arrays.
output_loss_metrics: List of metrics that are used to aggregated output
loss values.
sample_weights: Optional list of sample weight arrays.
training: The boolean represents if the weights of the model are updated.
'fit' methods will set this to True while 'evaluate' methods will
set this to False.
Returns:
output of the model, total loss, the loss and the mask
associated with each output.
Raises:
ValueError: If the model has no loss to optimize.
"""
with backend.eager_learning_phase_scope(1 if training else 0), \
training_utils.RespectCompiledTrainableState(model):
with GradientTape() as tape:
outs, total_loss, output_losses, masks = (_model_loss(
model,
inputs,
targets,
output_loss_metrics=output_loss_metrics,
sample_weights=sample_weights,
training=training))
if isinstance(model.optimizer,
loss_scale_optimizer.LossScaleOptimizer):
scaled_total_loss = model.optimizer.get_scaled_loss(total_loss)
else:
scaled_total_loss = total_loss
if training:
trainable_weights = model.trainable_weights
if trainable_weights:
# TODO(tanzheny) b/132690565: Provide mechanism for user to override
# model.train_on_batch.
if hasattr(model, '_backwards'):
model._backwards(tape, scaled_total_loss)
else:
grads = tape.gradient(scaled_total_loss, trainable_weights)
if isinstance(model.optimizer,
loss_scale_optimizer.LossScaleOptimizer):
grads = model.optimizer.get_unscaled_gradients(grads)
model.optimizer.apply_gradients(
zip(grads, trainable_weights))
else:
logging.warning(
'The list of trainable weights is empty. Make sure that'
' you are not setting model.trainable to False before '
'compiling the model.')
return outs, total_loss, output_losses, masks
def _process_single_batch(model,
inputs,
targets,
output_loss_metrics=None,
sample_weights=None,
training=False):
"""Calculate the loss and gradient for one input batch.
The model weights are updated if training is set to True.
Arguments:
model: Model whose loss has to be calculated.
inputs: List of input arrays.
targets: List of target arrays.
output_loss_metrics: List of metrics that are used to aggregated output
loss values.
sample_weights: Optional list of sample weight arrays.
training: The boolean represents if the weights of the model are updated.
'fit' methods will set this to True while 'evaluate' methods will
set this to False.
Returns:
output of the model, total loss, the loss and the mask
associated with each output.
Raises:
ValueError: If the model has no loss to optimize.
"""
with backend.eager_learning_phase_scope(1 if training else 0):
with GradientTape() as tape:
outs, total_loss, output_losses, masks = (_model_loss(
model,
inputs,
targets,
output_loss_metrics=output_loss_metrics,
sample_weights=sample_weights,
training=training))
if total_loss is None:
raise ValueError('The model cannot be run '
'because it has no loss to optimize.')
if isinstance(model.optimizer,
loss_scale_optimizer.LossScaleOptimizer):
# TODO(reedwm): Make loss_scale public instead of accessing private
# _loss_scale attribute.
loss_scale = model.optimizer._loss_scale()
scaled_total_loss = loss_scale_optimizer.scale_loss(
total_loss, loss_scale)
else:
loss_scale = None
scaled_total_loss = total_loss
if training:
if not model.trainable_weights:
logging.warning(
'The list of trainable weights is empty. Make sure that'
' you are not setting model.trainable to False before '
'compiling the model.')
else:
grads = tape.gradient(scaled_total_loss,
model.trainable_weights)
if loss_scale is not None:
grads = loss_scale_optimizer.unscale_grads(
grads, loss_scale)
model.optimizer.apply_gradients(
zip(grads, model.trainable_weights))
return outs, total_loss, output_losses, masks