def tpu_step(loss):
"""Generate the TPU graph."""
del loss
if is_training:
values = self.infeed_queue[0].generate_dequeue_op(tpu_device=0)
unflattened_inputs = data_nest.pack_sequence_as(
self.feature_structure, values)
else:
values = self.eval_infeed_queue[0].generate_dequeue_op(
tpu_device=0)
unflattened_inputs = data_nest.pack_sequence_as(
self.eval_feature_structure, values)
features = unflattened_inputs["features"]
labels = unflattened_inputs["labels"]
if is_training:
estimator_spec = model_fn(features, labels,
tf.estimator.ModeKeys.TRAIN, mparams)
loss, train_op = estimator_spec.loss, estimator_spec.train_op
with tf.device(device_for_tpu_core(self.get_host(0))):
with tf.control_dependencies([train_op]):
return tf.identity(loss)
else:
estimator_spec = model_fn(features, labels,
tf.estimator.ModeKeys.EVAL, mparams)
loss = estimator_spec.loss
self.eval_metrics = estimator_spec.eval_metrics
self.eval_tensors = estimator_spec.eval_metrics[1]
for _ in self.eval_tensors:
self.dequeue_ops.append([])
with tf.device(device_for_tpu_core(self.get_host(0))):
outfeed_enqueue_ops = tpu.outfeed_enqueue_tuple(
self.eval_tensors)
with tf.control_dependencies([outfeed_enqueue_ops]):
return tf.identity(loss)
def tpu_eval_step():
"""Generate the TPU graph."""
values = self.infeed_queue[0].generate_dequeue_op(tpu_device=0)
unflattened_inputs = data_nest.pack_sequence_as(self.feature_structure,
values)
features = unflattened_inputs["features"]
estimator_spec = model_fn(features, None, tf.estimator.ModeKeys.PREDICT,
params)
for k, v in six.iteritems(estimator_spec.predictions):
self.outfeed_names.append(k)
self.outfeed_tensors.append(v)
with tf.device(utils.device_for_tpu_core(self._get_host(0))):
outfeed_enqueue_ops = tpu.outfeed_enqueue_tuple(self.outfeed_tensors)
with tf.control_dependencies([outfeed_enqueue_ops]):
return tf.no_op()
def tpu_step(loss, *args):
"""Generate the TPU graph."""
del loss
unflattened_inputs = data_nest.pack_sequence_as(self.feature_structure,
args)
features = unflattened_inputs["features"]
labels = unflattened_inputs["labels"]
estimator_spec = model_fn(features, labels, tf.estimator.ModeKeys.EVAL,
mparams)
loss = estimator_spec.loss
self.eval_metrics = estimator_spec.eval_metrics
self.eval_tensors = estimator_spec.eval_metrics[1]
with tf.device(device_for_tpu_core()):
outfeed_enqueue_ops = tpu.outfeed_enqueue_tuple(self.eval_tensors)
with tf.control_dependencies([outfeed_enqueue_ops]):
return tf.identity(loss)