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_lib.ClassificationOutput(
constant_op.constant([[32.]]))})
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 test_classification_output_tensors_roundtrip_scores_only(self):
scores = array_ops.placeholder(dtypes.float32, 1, name='scores')
classification_output = export_lib.ClassificationOutput(scores=scores)
classification_output_tensors = (
tpu_estimator._export_output_to_tensors(classification_output))
self.assertEqual(classification_output_tensors, [scores, None])
scores_new = array_ops.placeholder(dtypes.float32,
1,
name='scores_new')
classification_output_new = (
tpu_estimator._clone_export_output_with_tensors(
classification_output, [scores_new, None]))
self.assertEqual(classification_output_new.scores, scores_new)
def test_classification_output_tensors_roundtrip_classes_only(self):
classes = array_ops.placeholder(dtypes.string, 1, name='classes')
classification_output = export_lib.ClassificationOutput(
classes=classes)
classification_output_tensors = (
tpu_estimator._export_output_to_tensors(classification_output))
self.assertEqual(classification_output_tensors, [None, classes])
classes_new = array_ops.placeholder(dtypes.string,
1,
name='classes_new')
classification_output_new = (
tpu_estimator._clone_export_output_with_tensors(
classification_output, [None, classes_new]))
self.assertEqual(classification_output_new.classes, classes_new)
def model_fn(features, labels, mode, params):
del params
loss = None
train_op = None
predictions = dense_computation(features)
export_outputs = None
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())
else:
if export_tpu_tensor:
key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
export_outputs = {
key: export_lib.PredictOutput({
'prediction': predictions
})
}
else:
export_outputs = {}
if export_cpu_tensor:
def host_call(predictions):
return string_ops.as_string(predictions, name='classes')
classes = tf.tpu.outside_compilation(host_call, predictions)
classification_output = export_lib.ClassificationOutput(
classes=classes)
export_outputs['classification'] = classification_output
if tpu_estimator_spec:
spec_type = tpu_estimator.TPUEstimatorSpec
else:
spec_type = model_fn_lib.EstimatorSpec
return spec_type(
mode,
loss=loss,
train_op=train_op,
predictions={'predictions': predictions},
export_outputs=export_outputs)
def test_classification_output_tensors_roundtrip_classify_both(self):
classes = array_ops.placeholder(dtypes.string, 1, name='classes')
scores = array_ops.placeholder(dtypes.float32, 1, name='scores')
classification_output = export_lib.ClassificationOutput(
scores, classes)
classification_output_tensors = (tpu_estimator.
_export_output_to_tensors(
classification_output))
self.assertSequenceEqual(classification_output_tensors, [scores, classes])
classes_new = array_ops.placeholder(dtypes.string, 1, name='classes_new')
scores_new = array_ops.placeholder(dtypes.float32, 1, name='scores_new')
classification_output_new = (tpu_estimator.
_clone_export_output_with_tensors(
classification_output,
[scores_new, classes_new]))
self.assertEqual(classification_output_new.classes, classes_new)
self.assertEqual(classification_output_new.scores, scores_new)