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, params):
del params
dense_features = tf.keras.layers.DenseFeatures(feature_columns)
input_layer = dense_features(features)
hidden_layer = tf.layers.dense(
input_layer,
HIDDEN_LAYER_SIZE,
kernel_initializer=tf.constant_initializer(KERNEL_INIT_VALUE),
bias_initializer=tf.constant_initializer(BIAS_INIT_VALUE))
last_layer = tf.reduce_sum(hidden_layer, axis=1)
logits = tf.reshape(last_layer, [-1])
labels = tf.reshape(labels, [-1])
losses = tf.square(labels - logits)
# Use reduce_mean to match the CrossShardOptimizer reduction.
loss = tf.reduce_mean(losses)
if optimizer_type == 'adagrad':
optimizer = tf.train.AdagradOptimizer(
LEARNING_RATE, initial_accumulator_value=ADADGRAD_INIT_VALUE)
elif optimizer_type == 'sgd':
optimizer = tf.train.GradientDescentOptimizer(LEARNING_RATE)
else:
raise ValueError('{} is not supported.'.format(optimizer_type))
# Default reduction=tf.losses.Reduction.MEAN
optimizer = tf.tpu.CrossShardOptimizer(optimizer)
train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op)
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(
tf.compat.v1.train.GradientDescentOptimizer(learning_rate=0.5))
train_op = optimizer.minimize(
loss, global_step=tf.compat.v1.train.get_global_step())
return tpu_estimator.TPUEstimatorSpec(mode=mode,
train_op=train_op,
loss=loss)
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_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 _create_estimator_spec(mode, loss=None, predictions=None,
export_outputs=None, eval_metrics=None,
train_op=None):
if use_tpu_estimator_spec:
return tpu_estimator.TPUEstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
predictions=predictions,
export_outputs=export_outputs,
eval_metrics=eval_metrics)
else:
return model_fn_lib.EstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
predictions=predictions,
export_outputs=export_outputs,
eval_metric_ops=(eval_metrics[0](*eval_metrics[1]) if eval_metrics
else None))
def _model_fn(
features: Dict[Text, tf.Tensor],
params: Dict[Text, int],
mode: model_fn_lib.ModeKeys,
) -> tpu_estimator.TPUEstimatorSpec:
"""A model which writes activations and sequence lengths to a file.
This method creates a model to extract the activations and sequence
lengths on each TPU core and pass them to a host call which writes them
to a file.
The model also applies an optimizer to the activations simply to avoid an
undefined gradients error.
Args:
features: A dictionary mapping keys to tensor inputs.
params: Parameters passed by TPUEstimator.
mode: Mode can be (TRAIN, EVAL, PREDICT).
Returns:
A TPUEstimatorSpec which holds the training_op that TPUEstimator will
run on TPU and the host_call that TPUEstimator will run on the host.
"""
del params
input_layer = tf.keras.experimental.SequenceFeatures([embedding_column])
activations, sequence_lengths = input_layer(features)
opt = tf.tpu.CrossShardOptimizer(tf.train.GradientDescentOptimizer(0.1))
loss = tf.reduce_sum(activations)
train_op = opt.minimize(loss, global_step=tf.train.get_global_step())
return tpu_estimator.TPUEstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
host_call=(_host_call, [activations, sequence_lengths]),
)