def model_fn(features, labels, mode):
_ = labels
step = training.get_global_step()
w = variable_scope.get_variable(
'w',
shape=[],
initializer=init_ops.zeros_initializer(),
dtype=dtypes.int64)
if estimator_lib.ModeKeys.TRAIN == mode:
# to consume features, we have control dependency
with ops.control_dependencies([features]):
step_inc = state_ops.assign_add(training.get_global_step(), 1)
with ops.control_dependencies([step_inc]):
assign_w_to_step_plus_2 = w.assign(step + 2)
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(3.),
train_op=assign_w_to_step_plus_2)
if estimator_lib.ModeKeys.EVAL == mode:
# to consume features, we have control dependency
with ops.control_dependencies([features]):
loss = constant_op.constant(5.)
return estimator_lib.EstimatorSpec(
mode,
loss=loss,
# w is constant in each step, so the mean.
# w = 0 if step==0 else step+2
eval_metric_ops={'mean_of_const': metrics_lib.mean(w)})
def model_fn(features, labels, mode):
_ = labels
step = training.get_global_step()
w = variable_scope.get_variable(
'w',
shape=[],
initializer=init_ops.zeros_initializer(),
dtype=dtypes.int64)
if estimator_lib.ModeKeys.TRAIN == mode:
# to consume features, we have control dependency
with ops.control_dependencies([features]):
step_inc = state_ops.assign_add(training.get_global_step(),
1)
with ops.control_dependencies([step_inc]):
assign_w_to_step_plus_2 = w.assign(step + 2)
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(3.),
train_op=assign_w_to_step_plus_2)
if estimator_lib.ModeKeys.EVAL == mode:
# to consume features, we have control dependency
with ops.control_dependencies([features]):
loss = constant_op.constant(5.)
return estimator_lib.EstimatorSpec(
mode,
loss=loss,
# w is constant in each step, so the mean.
# w = 0 if step==0 else step+2
eval_metric_ops={'mean_of_const': metrics_lib.mean(w)})
def model_fn_global_step_incrementer(features, labels, mode):
_, _ = features, labels
global_step = training.get_global_step()
return model_fn_lib.EstimatorSpec(
mode,
loss=constant_op.constant(1.),
train_op=state_ops.assign_add(global_step, 1))
def _create_global_step(self, graph):
"""Creates a global step suitable for TPUs.
Args:
graph: The graph in which to create the global step.
Returns:
A global step `Tensor`.
Raises:
ValueError: if the global step tensor is already defined.
"""
graph = graph or ops.get_default_graph()
if training.get_global_step(graph) is not None:
raise ValueError('"global_step" already exists.')
# Create in proper graph and base name_scope.
with graph.as_default() as g, g.name_scope(None):
return variable_scope.get_variable(
ops.GraphKeys.GLOBAL_STEP,
shape=[],
dtype=dtypes.int32,
initializer=init_ops.zeros_initializer(),
trainable=False,
use_resource=True,
collections=[
ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.GLOBAL_STEP
])
def _create_head_with_eval_metric_ops(self, mode, loss, eval_metric_ops):
"""Creates a head returning `TPUEstimatorSpec` based on mode.
This version contains eval that will not run on TPUs, where eval_metric_ops
has not been split into a metrics_fn that runs on CPUs. The intent is to
test the entire eval (model_fn forward pass) and metrics output on CPU.
Args:
mode: The mode such as TRAIN, EVAL.
loss: Training loss `Tensor`. Must be either scalar, or with shape `[1]`.
eval_metric_ops: Dict of metric results keyed by name.
Returns:
An EstimatorSpec for EVAL or TPUEstimatorSpec otherwise.
"""
if mode == _EVAL:
return model_fn_lib.EstimatorSpec(mode=mode,
eval_metric_ops=eval_metric_ops,
loss=loss)
# Train
optimizer = tf.compat.v1.tpu.CrossShardOptimizer(
training.GradientDescentOptimizer(learning_rate=0.5))
train_op = optimizer.minimize(loss,
global_step=training.get_global_step())
return tpu_estimator.TPUEstimatorSpec(mode=mode,
train_op=train_op,
loss=loss)
def _model_fn(features, labels, mode):
predictions = layers.dense(
features['x'], 1, kernel_initializer=init_ops.zeros_initializer())
export_outputs = {
'predictions': export_output.RegressionOutput(predictions)
}
if mode == model_fn_lib.ModeKeys.PREDICT:
return model_fn_lib.EstimatorSpec(
mode, predictions=predictions, export_outputs=export_outputs)
loss = losses.mean_squared_error(labels, predictions)
train_op = training.GradientDescentOptimizer(learning_rate=0.5).minimize(
loss, training.get_global_step())
eval_metric_ops = {
'absolute_error': metrics_lib.mean_absolute_error(
labels, predictions)
}
return model_fn_lib.EstimatorSpec(
mode,
predictions=predictions,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
export_outputs=export_outputs)
def _model_fn(features, labels, mode):
predictions = layers.dense(
features['x'], 1, kernel_initializer=init_ops.zeros_initializer())
export_outputs = {
'predictions': export.RegressionOutput(predictions)
}
if mode == model_fn_lib.ModeKeys.PREDICT:
return model_fn_lib.EstimatorSpec(
mode, predictions=predictions, export_outputs=export_outputs)
loss = losses.mean_squared_error(labels, predictions)
train_op = training.GradientDescentOptimizer(learning_rate=0.5).minimize(
loss, training.get_global_step())
eval_metric_ops = {
'absolute_error': metrics_lib.mean_absolute_error(
labels, predictions)
}
return model_fn_lib.EstimatorSpec(
mode,
predictions=predictions,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops,
export_outputs=export_outputs)
def model_fn_diff_modes(features, labels, mode):
_, _ = features, labels
v = variables.Variable(21, name='some_var')
train_op = None
loss = constant_op.constant(104)
if mode == model_fn_lib.ModeKeys.TRAIN:
loss = constant_op.constant(105)
predictions = constant_op.constant([501])
train_op = control_flow_ops.group(
state_ops.assign_add(training.get_global_step(), 1),
state_ops.assign_add(v, 3))
elif mode == model_fn_lib.ModeKeys.EVAL:
loss = constant_op.constant(106)
predictions = constant_op.constant([502])
else:
loss = constant_op.constant(107)
predictions = constant_op.constant([503])
return model_fn_lib.EstimatorSpec(
mode,
loss=loss,
train_op=train_op,
eval_metric_ops={
'abs_err': metrics_lib.mean_absolute_error(
constant_op.constant(0), predictions)},
predictions=predictions)
def model_fn_global_step_incrementer(features, labels, mode):
_, _ = features, labels
global_step = training.get_global_step()
return model_fn_lib.EstimatorSpec(
mode,
loss=constant_op.constant(1.),
train_op=state_ops.assign_add(global_step, 1))
def build_subnetwork(self,
features,
logits_dimension,
training,
iteration_step,
summary,
previous_ensemble=None):
assert features is not None
assert training is not None
assert iteration_step is not None
assert summary is not None
# Trainable variables collection should always be empty when
# build_subnetwork is called.
assert not tf_compat.v1.get_collection(
tf_compat.v1.GraphKeys.TRAINABLE_VARIABLES)
# Subnetworks get iteration steps instead of global steps.
step_name = "subnetwork_test/iteration_step"
assert step_name == tf_compat.tensor_name(
tf_compat.v1.train.get_global_step())
assert step_name == tf_compat.tensor_name(train.get_global_step())
assert step_name == tf_compat.tensor_name(training_util.get_global_step())
assert step_name == tf_compat.tensor_name(tf_v1.train.get_global_step())
assert step_name == tf_compat.tensor_name(
tf_compat.v1.train.get_or_create_global_step())
assert step_name == tf_compat.tensor_name(train.get_or_create_global_step())
assert step_name == tf_compat.tensor_name(
training_util.get_or_create_global_step())
assert step_name == tf_compat.tensor_name(
tf_v1.train.get_or_create_global_step())
# Subnetworks get scoped summaries.
assert "fake_scalar" == tf_compat.v1.summary.scalar("scalar", 1.)
assert "fake_image" == tf_compat.v1.summary.image("image", 1.)
assert "fake_histogram" == tf_compat.v1.summary.histogram("histogram", 1.)
assert "fake_audio" == tf_compat.v1.summary.audio("audio", 1., 1.)
last_layer = tu.dummy_tensor(shape=(2, 3))
def logits_fn(logits_dim):
return tf_compat.v1.layers.dense(
last_layer,
units=logits_dim,
kernel_initializer=tf_compat.v1.glorot_uniform_initializer(
seed=self._seed))
if self._multi_head:
logits = {
"head1": logits_fn(logits_dimension / 2),
"head2": logits_fn(logits_dimension / 2)
}
last_layer = {"head1": last_layer, "head2": last_layer}
else:
logits = logits_fn(logits_dimension)
return Subnetwork(
last_layer=logits if self._use_logits_last_layer else last_layer,
logits=logits,
complexity=2,
persisted_tensors={})
def _model_fn_train_only(features, labels):
v = variables.Variable(constant_op.constant(23), name='v')
return model_fn_lib.EstimatorSpec(ModeKeys.TRAIN,
predictions=features * labels,
loss=constant_op.constant(5) + v,
train_op=state_ops.assign_add(
training.get_global_step(), 1))
def model_fn_diff_modes(features, labels, mode):
_, _ = features, labels
v = variables.Variable(21, name='some_var')
train_op = None
loss = constant_op.constant(104)
if mode == model_fn_lib.ModeKeys.TRAIN:
loss = constant_op.constant(105)
predictions = constant_op.constant([501])
train_op = control_flow_ops.group(
state_ops.assign_add(training.get_global_step(), 1),
state_ops.assign_add(v, 3))
elif mode == model_fn_lib.ModeKeys.EVAL:
loss = constant_op.constant(106)
predictions = constant_op.constant([502])
else:
loss = constant_op.constant(107)
predictions = constant_op.constant([503])
return model_fn_lib.EstimatorSpec(mode,
loss=loss,
train_op=train_op,
eval_metric_ops={
'abs_err':
metrics_lib.mean_absolute_error(
constant_op.constant(0),
predictions)
},
predictions=predictions)
def model_fn(features, labels, mode, params):
del params # unused
with variable_scope.variable_scope('m', reuse=variable_scope.AUTO_REUSE):
w = variable_scope.get_variable('W', shape=[1000, 10])
logits = math_ops.matmul(features, w)
loss = losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
if mode == model_fn_lib.ModeKeys.TRAIN:
optimizer = training.RMSPropOptimizer(learning_rate=0.01)
optimizer = tpu_optimizer.CrossShardOptimizer(optimizer)
train_op = optimizer.minimize(loss, training.get_global_step())
return tpu_estimator.TPUEstimatorSpec(
mode=model_fn_lib.ModeKeys.TRAIN,
loss=loss,
train_op=train_op,
)
elif mode == model_fn_lib.ModeKeys.EVAL:
def metric_fn(labels, logits):
labels = math_ops.cast(labels, dtypes.int64)
logging.info('LABELS %s %s', labels, logits)
return {
'recall@1': metrics_lib.recall_at_k(labels, logits, 1),
'recall@5': metrics_lib.recall_at_k(labels, logits, 5),
}
loss = losses.sparse_softmax_cross_entropy(labels=labels,
logits=logits)
eval_metrics = (metric_fn, [labels, logits])
return tpu_estimator.TPUEstimatorSpec(mode=model_fn_lib.ModeKeys.EVAL,
loss=loss,
eval_metrics=eval_metrics)
def _model_fn_with_x_y(features, labels, mode):
_ = labels
variables.Variable(1., name='weight')
scores = constant_op.constant([3.])
classes = constant_op.constant(['wumpus'])
if mode == model_fn_lib.ModeKeys.PREDICT:
variables.Variable(36., name='name_collision')
return model_fn_lib.EstimatorSpec(
mode,
predictions=constant_op.constant(10.),
export_outputs={
'test': export_output.ClassificationOutput(scores, classes)
})
else:
prefix = 'eval_' if mode == model_fn_lib.ModeKeys.EVAL else ''
multiplied = math_ops.multiply(features['x'],
features['y'],
name='{}multiplied'.format(prefix))
metrics = {
'mean':
metrics_lib.mean(features['x'] - features['y'],
name='{}mean'.format(prefix))
}
variables.Variable(1., name='later_var')
variables.Variable(3., name='name_collision')
return model_fn_lib.EstimatorSpec(mode,
predictions=multiplied,
loss=constant_op.constant(1.),
train_op=state_ops.assign_add(
training.get_global_step(), 1),
eval_metric_ops=metrics)
def _create_global_step(self, graph):
"""Creates a global step suitable for TPUs.
Args:
graph: The graph in which to create the global step.
Returns:
A global step `Tensor`.
Raises:
ValueError: if the global step tensor is already defined.
"""
graph = graph or ops.get_default_graph()
if training.get_global_step(graph) is not None:
raise ValueError('"global_step" already exists.')
# Create in proper graph and base name_scope.
with graph.as_default() as g, g.name_scope(None):
return variable_scope.get_variable(
ops.GraphKeys.GLOBAL_STEP,
shape=[],
dtype=dtypes.int32,
initializer=init_ops.zeros_initializer(),
trainable=False,
use_resource=True,
collections=[ops.GraphKeys.GLOBAL_VARIABLES,
ops.GraphKeys.GLOBAL_STEP])
def model_fn(features, mode):
del features
global_step = training.get_global_step()
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant([5.]),
predictions={'x': constant_op.constant([5.])},
train_op=global_step.assign_add(1))
def model_fn(features, mode):
del features
global_step = training.get_global_step()
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant([5.]),
predictions={'x': constant_op.constant([5.])},
train_op=global_step.assign_add(1))
def model_fn(features, labels, mode):
_ = labels
with ops.control_dependencies([features['x']]):
loss = features['x'][1][0]
return model_fn_lib.EstimatorSpec(
mode,
loss=loss,
train_op=state_ops.assign_add(training.get_global_step(), 1))
def model_fn(features, labels, mode):
_ = labels
with ops.control_dependencies([features['x']]):
loss = features['x'][1][0]
return model_fn_lib.EstimatorSpec(
mode,
loss=loss,
train_op=state_ops.assign_add(training.get_global_step(), 1))
def model_fn(features, labels, mode):
_, _ = features, labels
return model_fn_lib.EstimatorSpec(
mode,
loss=constant_op.constant([103]),
train_op=state_ops.assign_add(training.get_global_step(), 1),
predictions=constant_op.constant([502]),
export_outputs={'test': export_output.ClassificationOutput(
constant_op.constant([[32.]]))})
def model_fn(features, labels, mode):
_, _ = features, labels
return model_fn_lib.EstimatorSpec(
mode,
loss=constant_op.constant([103]),
train_op=state_ops.assign_add(training.get_global_step(), 1),
predictions=constant_op.constant([502]),
export_outputs={'test': export_output.ClassificationOutput(
constant_op.constant([[32.]]))})
def _model_fn(features, labels, mode):
_, _, _ = features, labels, mode
self.assertIsNotNone(training.get_global_step())
self.assertEqual(expected_random_seed, ops.get_default_graph().seed)
return model_fn_lib.EstimatorSpec(
mode=mode,
loss=constant_op.constant(0.),
train_op=constant_op.constant(0.),
predictions=constant_op.constant([[0.]]))
def model_fn(features, labels, mode):
loss = None
if labels is not None:
loss = labels[0][0] + labels[1][0]
return model_fn_lib.EstimatorSpec(
mode,
loss=loss,
train_op=state_ops.assign_add(training.get_global_step(), 1),
predictions={'features_0': array_ops.identity([features['x'][0][0]]),
'features_1': array_ops.identity([features['x'][1][0]])})
def _model_fn_with_incremental_loss(features, labels, mode):
_, _ = features, labels
local_weight = variables.Variable(
0., name='local_weight', collections=[ops.GraphKeys.LOCAL_VARIABLES])
# Loss will be 2, 4, 6, ...
loss = 2 * state_ops.assign_add(local_weight, 1.)
return model_fn_lib.EstimatorSpec(
mode,
loss=loss,
train_op=state_ops.assign_add(training.get_global_step(), 1))
def _model_fn_with_incremental_loss(features, labels, mode):
_, _ = features, labels
local_weight = variables.Variable(
0., name='local_weight', collections=[ops.GraphKeys.LOCAL_VARIABLES])
# Loss will be 2, 4, 6, ...
loss = 2 * state_ops.assign_add(local_weight, 1.)
return model_fn_lib.EstimatorSpec(
mode,
loss=loss,
train_op=state_ops.assign_add(training.get_global_step(), 1))
def model_fn(features, labels, mode):
loss = None
if labels is not None:
loss = labels[0][0] + labels[1][0]
return model_fn_lib.EstimatorSpec(
mode,
loss=loss,
train_op=state_ops.assign_add(training.get_global_step(), 1),
predictions={'features_0': array_ops.identity([features['x'][0][0]]),
'features_1': array_ops.identity([features['x'][1][0]])})
def model_fn(features, labels, mode):
_, _ = features, labels
v = variables.Variable(21, name='some_var')
scaffold = monitored_session.Scaffold(
local_init_op=state_ops.assign_add(v, -3).op
)
return model_fn_lib.EstimatorSpec(
mode,
scaffold=scaffold,
train_op=state_ops.assign_add(training.get_global_step(), 1),
loss=array_ops.identity(v))
def _create_and_assert_global_step(self, graph):
"""Creates and asserts properties of the global step.
Args:
graph: The graph in which to create the global step tensor.
Returns:
The global step `Tensor`.
"""
step = self._create_global_step(graph)
assert step == training.get_global_step()
assert step.dtype.is_integer
return step
def _create_head(self, mode, loss, eval_metrics):
"""Creates a head returning `TPUEstimatorSpec` based on mode."""
if mode == _EVAL:
return tpu_estimator.TPUEstimatorSpec(
mode=mode, eval_metrics=eval_metrics, loss=loss)
# Train
optimizer = tf.tpu.CrossShardOptimizer(
training.GradientDescentOptimizer(learning_rate=0.5))
train_op = optimizer.minimize(loss, global_step=training.get_global_step())
return tpu_estimator.TPUEstimatorSpec(
mode=mode, train_op=train_op, loss=loss)
def model_fn(features, labels, mode):
_, _ = features, labels
v = variables.Variable(21, name='some_var')
scaffold = monitored_session.Scaffold(
local_init_op=state_ops.assign_add(v, -3).op)
return model_fn_lib.EstimatorSpec(mode,
scaffold=scaffold,
train_op=state_ops.assign_add(
training.get_global_step(),
1),
loss=array_ops.identity(v))
def model_fn(features, labels, mode):
_ = labels
if estimator_lib.ModeKeys.TRAIN == mode:
with ops.control_dependencies([features]):
train_op = state_ops.assign_add(training.get_global_step(), 1)
return estimator_lib.EstimatorSpec(
mode, loss=constant_op.constant(3.), train_op=train_op)
if estimator_lib.ModeKeys.EVAL == mode:
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(5.),
eval_metric_ops={'mean_of_features': metrics_lib.mean(features)})
def model_fn(features, labels, mode):
_ = labels
if estimator_lib.ModeKeys.TRAIN == mode:
with ops.control_dependencies([features]):
train_op = state_ops.assign_add(training.get_global_step(), 1)
return estimator_lib.EstimatorSpec(
mode, loss=constant_op.constant(3.), train_op=train_op)
if estimator_lib.ModeKeys.EVAL == mode:
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(5.),
eval_metric_ops={'mean_of_features': metrics_lib.mean(features)})
def _create_and_assert_global_step(self, graph):
"""Creates and asserts properties of the global step.
Args:
graph: The graph in which to create the global step tensor.
Returns:
The global step `Tensor`.
"""
step = self._create_global_step(graph)
assert step == training.get_global_step()
assert step.dtype.is_integer
return step
def model_fn(features, labels, mode):
tb = lookup_ops.MutableHashTable(key_dtype=dtypes.int32,
value_dtype=dtypes.int32,
default_value=-1)
predictions = tb.lookup(features['x'])
train_op = None
if mode == ModeKeys.TRAIN:
train_op = control_flow_ops.group(
tb.insert(features['x'], labels),
state_ops.assign_add(training.get_global_step(), 1))
return model_fn_lib.EstimatorSpec(mode,
loss=constant_op.constant(0),
predictions=predictions,
train_op=train_op)
def model_fn(features, labels, mode, params):
loss = None
train_op = None
export_outputs = None
# This could be some pre-processing on CPU like calls to input layer with
# embedding columns.
x2 = features['x'] * 2
def computation(input_tensor):
return layers.dense(
input_tensor,
1,
kernel_initializer=init_ops.zeros_initializer())
if mode != _PREDICT:
predictions = computation(x2)
loss = losses.mean_squared_error(labels, predictions)
optimizer = tf.tpu.CrossShardOptimizer(
training.GradientDescentOptimizer(learning_rate=0.5))
train_op = optimizer.minimize(loss, training.get_global_step())
else:
inputs = [x2]
if params['use_tpu']:
predictions = array_ops.identity(
tpu_estimator.inference_on_tpu(computation,
inputs,
num_batch_threads=1,
max_batch_size=2,
batch_timeout_micros=100),
name='predictions')
else:
predictions = array_ops.identity(computation(*inputs),
name='predictions')
key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
export_outputs = {
key: export_lib.PredictOutput({'prediction': predictions})
}
classes = string_ops.as_string(predictions, name='classes')
classification_output = export_lib.ClassificationOutput(
classes=classes)
export_outputs['classification'] = classification_output
return tpu_estimator.TPUEstimatorSpec(
mode,
loss=loss,
train_op=train_op,
predictions={'predictions': predictions},
export_outputs=export_outputs)
def _model_fn(features, labels, mode, config, params=None):
"""model_fn."""
# TODO(jhseu): Move to EVAL and PREDICT to TPU.
if mode != model_fn_lib.ModeKeys.TRAIN:
return _call_model_fn_without_tpu(model_fn, features, labels, mode,
config, params)
# Now for TPU training.
if params is not None and _BATCH_SIZE_KEY in params:
params[_BATCH_SIZE_KEY] //= config.tpu_config.num_shards
assert isinstance(features, _PerShardOutput)
features = features.as_list()
if labels is not None:
assert isinstance(labels, _PerShardOutput)
labels = labels.as_list()
dequeue_fn, enqueue_fn = (_create_infeed_enqueue_ops_and_dequeue_fn(
config, features, labels))
loss = _train_on_tpu_shards(config,
train_step=_convert_model_fn_to_train_step(
model_fn, dequeue_fn, mode, config,
params))
# Gets the variables back from TPU nodes. This means the variables updated
# by TPU will now be *synced* to host memory.
update_ops = [
array_ops.check_numerics(v.read_value(),
'Gradient for %s is NaN' % v.name).op
for v in variables.trainable_variables()
]
hooks = [
TpuInfeedSessionHook(config, enqueue_fn),
training.LoggingTensorHook(
{
'loss': array_ops.identity(loss),
'step': training.get_global_step()
},
every_n_secs=30)
]
return model_fn_lib.EstimatorSpec(
mode,
loss=array_ops.identity(loss),
training_hooks=hooks,
train_op=control_flow_ops.group(*update_ops))
def _model_fn(features, labels, mode):
v = variables.Variable(constant_op.constant(23), name='v')
if mode == ModeKeys.PREDICT:
return model_fn_lib.EstimatorSpec(ModeKeys.PREDICT,
predictions=features + 1)
elif mode == ModeKeys.EVAL:
return model_fn_lib.EstimatorSpec(ModeKeys.EVAL,
loss=constant_op.constant(5) + v,
predictions=features + labels)
elif mode == ModeKeys.TRAIN:
return model_fn_lib.EstimatorSpec(ModeKeys.TRAIN,
predictions=features * labels,
loss=constant_op.constant(5) + v,
train_op=state_ops.assign_add(
training.get_global_step(), 1))
def _model_fn_with_eval_metric_ops(features, labels, mode, params):
_, _ = features, labels
metric_name = params.get('metric_name') or 'metric'
metric_value = params.get('metric_value') or 2.
global_step = training.get_global_step()
loss = constant_op.constant(1.)
metric_update_op = loss.op
metric_tensor = control_flow_ops.with_dependencies(
[metric_update_op], constant_op.constant(metric_value))
return model_fn_lib.EstimatorSpec(
mode,
loss=loss,
predictions={'predictions': constant_op.constant(1.)},
train_op=state_ops.assign_add(global_step, 1),
eval_metric_ops={metric_name: (metric_tensor, metric_update_op)})
def _model_fn_with_eval_metric_ops(features, labels, mode, params):
_, _ = features, labels
metric_name = params.get('metric_name') or 'metric'
metric_value = params.get('metric_value') or 2.
global_step = training.get_global_step()
loss = constant_op.constant(1.)
metric_update_op = loss.op
metric_tensor = control_flow_ops.with_dependencies(
[metric_update_op], constant_op.constant(metric_value))
return model_fn_lib.EstimatorSpec(
mode,
loss=loss,
predictions={'predictions': constant_op.constant(1.)},
train_op=state_ops.assign_add(global_step, 1),
eval_metric_ops={metric_name: (metric_tensor, metric_update_op)})
def _model_fn_for_export_tests(features, labels, mode):
_, _ = features, labels
variables.Variable(1., name='weight')
scores = constant_op.constant([3.])
classes = constant_op.constant(['wumpus'])
update_global_step = state_ops.assign_add(training.get_global_step(), 1)
with ops.control_dependencies([update_global_step]):
train_op = constant_op.constant(2.)
return model_fn_lib.EstimatorSpec(
mode,
predictions=constant_op.constant(10.),
loss=constant_op.constant(1.),
train_op=train_op,
export_outputs={
'test': export_output.ClassificationOutput(scores, classes)})
def model_fn(features, labels, mode):
_, _ = features, labels
w = variables.Variable(
initial_value=[0.],
trainable=False,
collections=[ops.GraphKeys.SAVEABLE_OBJECTS])
init_op = control_flow_ops.group(
[w.initializer, training.get_global_step().initializer])
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(3.),
scaffold=training.Scaffold(init_op=init_op),
train_op=constant_op.constant(5.),
eval_metric_ops={
'mean_of_features': metrics_lib.mean(constant_op.constant(2.))
})
def model_fn(features, labels, mode):
_, _ = features, labels
w = variables.VariableV1(
initial_value=[0.],
trainable=False,
collections=[ops.GraphKeys.SAVEABLE_OBJECTS])
init_op = control_flow_ops.group(
[w.initializer, training.get_global_step().initializer])
return estimator_lib.EstimatorSpec(
mode,
loss=constant_op.constant(3.),
scaffold=training.Scaffold(init_op=init_op),
train_op=constant_op.constant(5.),
eval_metric_ops={
'mean_of_features': metrics_lib.mean(constant_op.constant(2.))
})
def _model_fn(features, labels, mode, config, params=None):
"""model_fn."""
# TODO(jhseu): Move to EVAL and PREDICT to TPU.
if mode != model_fn_lib.ModeKeys.TRAIN:
return _call_model_fn_without_tpu(
model_fn, features, labels, mode, config, params)
# Now for TPU training.
if params is not None and _BATCH_SIZE_KEY in params:
params[_BATCH_SIZE_KEY] //= config.tpu_config.num_shards
assert isinstance(features, _PerShardOutput)
features = features.as_list()
if labels is not None:
assert isinstance(labels, _PerShardOutput)
labels = labels.as_list()
dequeue_fn, enqueue_fn = (
_create_infeed_enqueue_ops_and_dequeue_fn(config, features, labels))
loss = _train_on_tpu_shards(
config,
train_step=_convert_model_fn_to_train_step(
model_fn, dequeue_fn, mode, config, params))
# Gets the variables back from TPU nodes. This means the variables updated
# by TPU will now be *synced* to host memory.
update_ops = [
array_ops.check_numerics(v.read_value(),
'Gradient for %s is NaN' % v.name).op
for v in variables.trainable_variables()
]
hooks = [
TpuInfeedSessionHook(config, enqueue_fn),
training.LoggingTensorHook(
{'loss': array_ops.identity(loss),
'step': training.get_global_step()},
every_n_secs=30)
]
return model_fn_lib.EstimatorSpec(
mode,
loss=array_ops.identity(loss),
training_hooks=hooks,
train_op=control_flow_ops.group(*update_ops))
def _model_fn(features, labels, mode, params):
if not self._export_mode:
# Always check batch size in params
self.assertEqual(batch_size_dict[mode], params['batch_size'])
else:
self.assertNotIn('batch_size', params)
# Check the input feeds correct shape for train and eval. When eval on CPU
# or predict, it is allowed to have dynamic shape. So, here only validates
# the fully known shape (which covers the TPU train).
if features['x'].shape.is_fully_defined():
self.assertEqual(batch_size_dict[mode], features['x'].shape[0])
predictions = layers.dense(
features['x'],
1,
kernel_initializer=init_ops.ones_initializer())
export_outputs = {
'predictions': export_output.RegressionOutput(predictions)
}
if mode == _PREDICT:
return _create_estimator_spec(
mode=mode,
predictions={'predictions': predictions},
export_outputs=export_outputs)
loss = losses.mean_squared_error(labels, predictions)
optimizer = tf.tpu.CrossShardOptimizer(
training.GradientDescentOptimizer(learning_rate=0.5))
train_op = optimizer.minimize(
loss, global_step=training.get_global_step())
eval_metrics = (
lambda labels, predictions: { # pylint: disable=g-long-lambda
'absolute_error':
metrics_lib.mean_absolute_error(labels, predictions)
},
[labels, predictions])
return _create_estimator_spec(
mode=mode,
loss=loss,
predictions={'predictions': predictions},
export_outputs=export_outputs,
train_op=train_op,
eval_metrics=eval_metrics)
def _model_fn_callable_variable_initializers(features, labels, mode):
"""Model_fn with callable variable initializers (for WrappedGraph tests)."""
_ = features, labels
v = variables.Variable(lambda: constant_op.constant(23), name='v')
if mode == ModeKeys.PREDICT:
return model_fn_lib.EstimatorSpec(ModeKeys.PREDICT,
predictions=features + 1)
elif mode == ModeKeys.EVAL:
return model_fn_lib.EstimatorSpec(ModeKeys.EVAL,
loss=constant_op.constant(5) + v,
predictions=features + labels)
elif mode == ModeKeys.TRAIN:
return model_fn_lib.EstimatorSpec(ModeKeys.TRAIN,
predictions=features * labels,
loss=constant_op.constant(5) + v,
train_op=state_ops.assign_add(
training.get_global_step(), 1))
def model_fn_global_step_incrementer(features, labels, mode, params):
del params
loss = None
train_op = None
predictions = dense_computation(features)
if mode != _PREDICT:
loss = losses.mean_squared_error(labels, predictions)
optimizer = tf.tpu.CrossShardOptimizer(
training.GradientDescentOptimizer(learning_rate=0.5))
train_op = optimizer.minimize(loss, training.get_global_step())
return tpu_estimator.TPUEstimatorSpec(
mode,
loss=loss,
train_op=train_op,
predictions={'predictions': predictions},
export_outputs={
'test': export_output.PredictOutput({'prediction': predictions})
})
def _build_train_op(self, loss):
"""Creates the training operation,
In case of use_target_network == True, we append also the update op
while taking into account the update_frequency.
"""
train_op = super(BaseQModel, self)._build_train_op(loss)
# check if we need to update the target graph
if self.use_target_graph:
update_op = tf.cond(
tf.equal(tf.mod(training.get_global_step(), self.target_update_frequency), 0),
self._build_update_target_graph,
lambda: tf.no_op(name='no_op_copy_target'))
# append the target update op to the train op.
train_op = tf.group(*[train_op, update_op], name='train_and_update_target')
return train_op
def _model_fn(features, labels, mode, config, params):
"""A Estimator `model_fn` for TPUEstimator."""
model_fn_wrapper = _ModelFnWrapper(model_fn, config, params, mode,
train_batch_size)
# TODO(jhseu): Move to EVAL and PREDICT to TPU.
if not use_tpu or mode != model_fn_lib.ModeKeys.TRAIN:
return model_fn_wrapper.call_without_tpu(features, labels)
inputs = _InputsHolder(features=features, labels=labels,
num_shards=config.tpu_config.num_shards)
dequeue_fn, enqueue_fn = _create_infeed_enqueue_ops_and_dequeue_fn(
inputs, config)
loss = _train_on_tpu_system(model_fn_wrapper, dequeue_fn)
# Gets the variables back from TPU nodes. This means the variables updated
# by TPU will now be *synced* to host memory.
update_ops = [
array_ops.check_numerics(v.read_value(),
'Gradient for %s is NaN' % v.name).op
for v in variables.trainable_variables()
]
hooks = [
TPUInfeedSessionHook(config, enqueue_fn),
training.LoggingTensorHook(
{'loss': array_ops.identity(loss),
'step': training.get_global_step()},
every_n_secs=30)
]
return model_fn_lib.EstimatorSpec(
mode,
loss=array_ops.identity(loss),
training_hooks=hooks,
train_op=control_flow_ops.group(*update_ops))
def get_grad_multiplier(self):
if self._grad_multiplier_fn:
return ops.convert_to_tensor(
self._grad_multiplier_fn(training.get_global_step()),
dtype=dtypes.float32)
def optimize_loss(loss,
global_step,
learning_rate,
optimizer,
gradient_noise_scale=None,
gradient_multipliers=None,
clip_gradients=None,
learning_rate_decay_fn=None,
update_ops=None,
variables=None,
name=None,
summaries=None,
colocate_gradients_with_ops=False,
increment_global_step=True):
"""Given loss and parameters for optimizer, returns a training op.
Various ways of passing optimizers include:
- by string specifying the name of the optimizer. See OPTIMIZER_CLS_NAMES
for full list. E.g. `optimize_loss(..., optimizer='Adam')`.
- by function taking learning rate `Tensor` as argument and returning an
`Optimizer` instance. E.g. `optimize_loss(...,
optimizer=lambda lr: tf.train.MomentumOptimizer(lr, momentum=0.5))`.
Alternatively, if `learning_rate` is `None`, the function takes no
arguments. E.g. `optimize_loss(..., learning_rate=None,
optimizer=lambda: tf.train.MomentumOptimizer(0.5, momentum=0.5))`.
- by a subclass of `Optimizer` having a single-argument constructor
(the argument is the learning rate), such as AdamOptimizer or
AdagradOptimizer. E.g. `optimize_loss(...,
optimizer=tf.train.AdagradOptimizer)`.
- by an instance of a subclass of `Optimizer`.
E.g., `optimize_loss(..., optimizer=tf.train.AdagradOptimizer(0.5))`.
Args:
loss: Scalar `Tensor`.
global_step: Scalar int `Tensor`, step counter to update on each step
unless `increment_global_step` is `False`. If not supplied,
it will be fetched from the default graph (see
`tf.train.get_global_step` for details). If it has
not been created, no step will be incremented with each weight
update. `learning_rate_decay_fn` requires `global_step`.
learning_rate: float or `Tensor`, magnitude of update per each training
step. Can be `None`.
optimizer: string, class or optimizer instance, used as trainer.
string should be name of optimizer, like 'SGD',
'Adam', 'Adagrad'. Full list in OPTIMIZER_CLS_NAMES constant.
class should be sub-class of `tf.Optimizer` that implements
`compute_gradients` and `apply_gradients` functions.
optimizer instance should be instantiation of `tf.Optimizer`
sub-class and have `compute_gradients` and `apply_gradients`
functions.
gradient_noise_scale: float or None, adds 0-mean normal noise scaled by this
value.
gradient_multipliers: dict of variables or variable names to floats.
If present, gradients for specified
variables will be multiplied by given constant.
clip_gradients: float, callable or `None`. If float, is provided, a global
clipping is applied to prevent the norm of the gradient to exceed this
value. Alternatively, a callable can be provided e.g.: adaptive_clipping.
This callable takes a `list` of `(gradients, variables)` `tuple`s and
returns the same thing with the gradients modified.
learning_rate_decay_fn: function, takes `learning_rate` and `global_step`
`Tensor`s, returns `Tensor`.
Can be used to implement any learning rate decay
functions.
For example: `tf.train.exponential_decay`.
Ignored if `learning_rate` is not supplied.
update_ops: list of update `Operation`s to execute at each step. If `None`,
uses elements of UPDATE_OPS collection. The order of execution
between `update_ops` and `loss` is non-deterministic.
variables: list of variables to optimize or
`None` to use all trainable variables.
name: The name for this operation is used to scope operations and summaries.
summaries: List of internal quantities to visualize on tensorboard. If not
set, the loss, the learning rate, and the global norm of the
gradients will be reported. The complete list of possible values
is in OPTIMIZER_SUMMARIES.
colocate_gradients_with_ops: If True, try colocating gradients with the
corresponding op.
increment_global_step: Whether to increment `global_step`. If your model
calls `optimize_loss` multiple times per training step (e.g. to optimize
different parts of the model), use this arg to avoid incrementing
`global_step` more times than necessary.
Returns:
Training op.
Raises:
ValueError: if:
* `loss` is an invalid type or shape.
* `global_step` is an invalid type or shape.
* `learning_rate` is an invalid type or value.
* `optimizer` has the wrong type.
* `clip_gradients` is neither float nor callable.
* `learning_rate` and `learning_rate_decay_fn` are supplied, but no
`global_step` is available.
* `gradients` is empty.
"""
loss = ops.convert_to_tensor(loss)
contrib_framework.assert_scalar(loss)
if global_step is None:
global_step = train.get_global_step()
else:
train.assert_global_step(global_step)
with vs.variable_scope(name, "OptimizeLoss", [loss, global_step]):
# Update ops take UPDATE_OPS collection if not provided.
if update_ops is None:
update_ops = set(ops.get_collection(ops.GraphKeys.UPDATE_OPS))
# Make sure update ops are ran before computing loss.
if update_ops:
loss = control_flow_ops.with_dependencies(list(update_ops), loss)
# Learning rate variable, with possible decay.
lr = None
if learning_rate is not None:
if (isinstance(learning_rate, ops.Tensor) and
learning_rate.get_shape().ndims == 0):
lr = learning_rate
elif isinstance(learning_rate, float):
if learning_rate < 0.0:
raise ValueError("Invalid learning_rate %s.", learning_rate)
lr = vs.get_variable(
"learning_rate", [],
trainable=False,
initializer=init_ops.constant_initializer(learning_rate))
else:
raise ValueError("Learning rate should be 0d Tensor or float. "
"Got %s of type %s" % (str(learning_rate),
str(type(learning_rate))))
if summaries is None:
summaries = ["loss", "learning_rate", "global_gradient_norm"]
else:
for summ in summaries:
if summ not in OPTIMIZER_SUMMARIES:
raise ValueError("Summaries should be one of [%s], you provided %s." %
(", ".join(OPTIMIZER_SUMMARIES), summ))
if learning_rate is not None and learning_rate_decay_fn is not None:
if global_step is None:
raise ValueError("global_step is required for learning_rate_decay_fn.")
lr = learning_rate_decay_fn(lr, global_step)
if "learning_rate" in summaries:
summary.scalar("learning_rate", lr)
# Create optimizer, given specified parameters.
if isinstance(optimizer, six.string_types):
if lr is None:
raise ValueError("Learning rate is None, but should be specified if "
"optimizer is string (%s)." % optimizer)
if optimizer not in OPTIMIZER_CLS_NAMES:
raise ValueError(
"Optimizer name should be one of [%s], you provided %s." %
(", ".join(OPTIMIZER_CLS_NAMES), optimizer))
opt = OPTIMIZER_CLS_NAMES[optimizer](learning_rate=lr)
elif (isinstance(optimizer, type) and
issubclass(optimizer, optimizer_.Optimizer)):
if lr is None:
raise ValueError("Learning rate is None, but should be specified if "
"optimizer is class (%s)." % optimizer)
opt = optimizer(learning_rate=lr)
elif isinstance(optimizer, optimizer_.Optimizer):
opt = optimizer
elif callable(optimizer):
if learning_rate is not None:
opt = optimizer(lr)
else:
opt = optimizer()
if not isinstance(opt, optimizer_.Optimizer):
raise ValueError("Unrecognized optimizer: function should return "
"subclass of Optimizer. Got %s." % str(opt))
else:
raise ValueError("Unrecognized optimizer: should be string, "
"subclass of Optimizer, instance of "
"subclass of Optimizer or function with one argument. "
"Got %s." % str(optimizer))
# All trainable variables, if specific variables are not specified.
if variables is None:
variables = vars_.trainable_variables()
# Compute gradients.
gradients = opt.compute_gradients(
loss,
variables,
colocate_gradients_with_ops=colocate_gradients_with_ops)
# Optionally add gradient noise.
if gradient_noise_scale is not None:
gradients = _add_scaled_noise_to_gradients(gradients,
gradient_noise_scale)
# Multiply some gradients.
if gradient_multipliers is not None:
gradients = _multiply_gradients(gradients, gradient_multipliers)
if not gradients:
raise ValueError(
"Empty list of (gradient, var) pairs encountered. This is most "
"likely to be caused by an improper value of gradient_multipliers.")
if "global_gradient_norm" in summaries or "gradient_norm" in summaries:
summary.scalar("global_norm/gradient_norm",
clip_ops.global_norm(list(zip(*gradients))[0]))
# Optionally clip gradients by global norm.
if isinstance(clip_gradients, float):
gradients = _clip_gradients_by_norm(gradients, clip_gradients)
elif callable(clip_gradients):
gradients = clip_gradients(gradients)
elif clip_gradients is not None:
raise ValueError(
"Unknown type %s for clip_gradients" % type(clip_gradients))
# Add scalar summary for loss.
if "loss" in summaries:
summary.scalar("loss", loss)
# Add histograms for variables, gradients and gradient norms.
for gradient, variable in gradients:
if isinstance(gradient, ops.IndexedSlices):
grad_values = gradient.values
else:
grad_values = gradient
if grad_values is not None:
var_name = variable.name.replace(":", "_")
if "gradients" in summaries:
summary.histogram("gradients/%s" % var_name, grad_values)
if "gradient_norm" in summaries:
summary.scalar("gradient_norm/%s" % var_name,
clip_ops.global_norm([grad_values]))
if clip_gradients is not None and ("global_gradient_norm" in summaries or
"gradient_norm" in summaries):
summary.scalar("global_norm/clipped_gradient_norm",
clip_ops.global_norm(list(zip(*gradients))[0]))
# Create gradient updates.
grad_updates = opt.apply_gradients(
gradients,
global_step=global_step if increment_global_step else None,
name="train")
# Ensure the train_tensor computes grad_updates.
train_tensor = control_flow_ops.with_dependencies([grad_updates], loss)
return train_tensor
def testSaveAndLoadSavedModelExport(
self, model_builder, uses_learning_phase, optimizer, train_before_export):
saved_model_path = self._save_model_dir()
with self.session(graph=ops.Graph()):
np.random.seed(130)
input_arr = np.random.random((1, 3))
target_arr = np.random.random((1, 3))
model = model_builder(uses_learning_phase)
if optimizer is not None:
model.compile(
loss='mse',
optimizer=optimizer,
metrics=['mae'])
if train_before_export:
model.train_on_batch(input_arr, target_arr)
ref_loss, ref_mae = model.evaluate(input_arr, target_arr)
ref_predict = model.predict(input_arr)
# Export SavedModel
output_path = keras_saved_model.save_keras_model(model, saved_model_path)
input_name = model.input_names[0]
output_name = model.output_names[0]
target_name = output_name + '_target'
# Load predict graph, and test predictions
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs, _ = load_model(sess, output_path,
model_fn_lib.ModeKeys.PREDICT)
predictions = sess.run(outputs[output_name],
{inputs[input_name]: input_arr})
self.assertAllClose(ref_predict, predictions, atol=1e-05)
if optimizer:
# Load eval graph, and test predictions, loss and metric values
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs, _ = load_model(sess, output_path,
model_fn_lib.ModeKeys.EVAL)
# First obtain the loss and predictions, and run the metric update op by
# feeding in the inputs and targets.
loss, predictions, _ = sess.run(
(outputs['loss'], outputs['predictions/' + output_name],
outputs['metrics/mean_absolute_error/update_op']), {
inputs[input_name]: input_arr,
inputs[target_name]: target_arr
})
# The metric value should be run after the update op, to ensure that it
# reflects the correct value.
metric_value = sess.run(outputs['metrics/mean_absolute_error/value'])
self.assertEqual(int(train_before_export),
sess.run(training_module.get_global_step()))
self.assertAllClose(ref_loss, loss, atol=1e-05)
self.assertAllClose(ref_mae, metric_value, atol=1e-05)
self.assertAllClose(ref_predict, predictions, atol=1e-05)
# Load train graph, and check for the train op, and prediction values
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs, meta_graph_def = load_model(
sess, output_path, model_fn_lib.ModeKeys.TRAIN)
self.assertEqual(int(train_before_export),
sess.run(training_module.get_global_step()))
self.assertIn('loss', outputs)
self.assertIn('metrics/mean_absolute_error/update_op', outputs)
self.assertIn('metrics/mean_absolute_error/value', outputs)
self.assertIn('predictions/' + output_name, outputs)
# Train for a step
train_op = loader_impl.get_train_op(meta_graph_def)
train_outputs, _ = sess.run(
[outputs, train_op], {inputs[input_name]: input_arr,
inputs[target_name]: target_arr})
self.assertEqual(int(train_before_export) + 1,
sess.run(training_module.get_global_step()))
if uses_learning_phase:
self.assertAllClose(
[[0, 0, 0]], train_outputs['predictions/' + output_name],
atol=1e-05)
else:
self.assertNotAllClose(
[[0, 0, 0]], train_outputs['predictions/' + output_name],
atol=1e-05)
def testSaveAndLoadSavedModelExport(
self, model_builder, uses_learning_phase, optimizer, train_before_export):
saved_model_path = self._save_model_dir()
with self.session(graph=ops.Graph()):
input_arr = np.random.random((1, 3))
target_arr = np.random.random((1, 3))
model = model_builder(uses_learning_phase)
if optimizer is not None:
model.compile(
loss='mse',
optimizer=optimizer,
metrics=['mae'])
if train_before_export:
model.train_on_batch(input_arr, target_arr)
ref_loss, ref_mae = model.evaluate(input_arr, target_arr)
ref_predict = model.predict(input_arr)
# Export SavedModel
output_path = keras_saved_model.save_keras_model(model, saved_model_path)
input_name = model.input_names[0]
output_name = model.output_names[0]
target_name = output_name + '_target'
# Load predict graph, and test predictions
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs = load_model(sess, output_path,
model_fn_lib.ModeKeys.PREDICT)
predictions = sess.run(outputs[output_name],
{inputs[input_name]: input_arr})
self.assertAllClose(ref_predict, predictions, atol=1e-05)
if optimizer:
# Load eval graph, and test predictions, loss and metric values
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs = load_model(sess, output_path,
model_fn_lib.ModeKeys.EVAL)
eval_results = sess.run(outputs, {inputs[input_name]: input_arr,
inputs[target_name]: target_arr})
self.assertEqual(int(train_before_export),
sess.run(training_module.get_global_step()))
self.assertAllClose(ref_loss, eval_results['loss'], atol=1e-05)
self.assertAllClose(
ref_mae, eval_results['metrics/mae/update_op'], atol=1e-05)
self.assertAllClose(
ref_predict, eval_results['predictions/' + output_name], atol=1e-05)
# Load train graph, and check for the train op, and prediction values
with session.Session(graph=ops.Graph()) as sess:
inputs, outputs = load_model(sess, output_path,
model_fn_lib.ModeKeys.TRAIN)
self.assertEqual(int(train_before_export),
sess.run(training_module.get_global_step()))
self.assertIn('loss', outputs)
self.assertIn('metrics/mae/update_op', outputs)
self.assertIn('metrics/mae/value', outputs)
self.assertIn('predictions/' + output_name, outputs)
# Train for a step
train_op = ops.get_collection(constants.TRAIN_OP_KEY)
train_outputs, _ = sess.run(
[outputs, train_op], {inputs[input_name]: input_arr,
inputs[target_name]: target_arr})
self.assertEqual(int(train_before_export) + 1,
sess.run(training_module.get_global_step()))
if uses_learning_phase:
self.assertAllClose(
[[0, 0, 0]], train_outputs['predictions/' + output_name],
atol=1e-05)
else:
self.assertNotAllClose(
[[0, 0, 0]], train_outputs['predictions/' + output_name],
atol=1e-05)