def _serving_ops(self, features):
"""Add ops for serving to the graph."""
with variable_scope.variable_scope("model", use_resource=True):
filtering_features = {}
prediction_features = {}
values_length = array_ops.shape(
features[feature_keys.FilteringFeatures.VALUES])[1]
for key, value in features.items():
if key == feature_keys.State.STATE_TUPLE:
# Ignore state input. The model's default start state is replicated
# across the batch.
continue
if key == feature_keys.FilteringFeatures.VALUES:
filtering_features[key] = value
else:
filtering_features[key] = value[:, :values_length]
prediction_features[key] = value[:, values_length:]
cold_filtering_outputs = self.model.define_loss(
features=filtering_features, mode=estimator_lib.ModeKeys.EVAL)
prediction_features[feature_keys.State.STATE_TUPLE] = (
cold_filtering_outputs.end_state)
with variable_scope.variable_scope("model", reuse=True):
prediction_outputs = self.model.predict(
features=prediction_features)
return estimator_lib.EstimatorSpec(
mode=estimator_lib.ModeKeys.PREDICT,
export_outputs={
feature_keys.SavedModelLabels.PREDICT:
_NoStatePredictOutput(prediction_outputs),
},
# Likely unused, but it is necessary to return `predictions` to satisfy
# the Estimator's error checking.
predictions={})
def _serving_ops(self, features):
"""Add ops for serving to the graph."""
with variable_scope.variable_scope("model", use_resource=True):
prediction_outputs = self.model.predict(features=features)
with variable_scope.variable_scope("model", reuse=True):
filtering_outputs = self.create_loss(features,
estimator_lib.ModeKeys.EVAL)
with variable_scope.variable_scope("model", reuse=True):
no_state_features = {
k: v
for k, v in features.items()
if not k.startswith(feature_keys.State.STATE_PREFIX)
}
# Ignore any state management when cold-starting. The model's default
# start state is replicated across the batch.
cold_filtering_outputs = self.model.define_loss(
features=no_state_features, mode=estimator_lib.ModeKeys.EVAL)
return estimator_lib.EstimatorSpec(
mode=estimator_lib.ModeKeys.PREDICT,
export_outputs={
feature_keys.SavedModelLabels.PREDICT:
export_lib.PredictOutput(prediction_outputs),
feature_keys.SavedModelLabels.FILTER:
export_lib.PredictOutput(
state_to_dictionary(filtering_outputs.end_state)),
feature_keys.SavedModelLabels.COLD_START_FILTER:
_NoStatePredictOutput(
state_to_dictionary(cold_filtering_outputs.end_state))
},
# Likely unused, but it is necessary to return `predictions` to satisfy
# the Estimator's error checking.
predictions={})
def _evaluate_ops(self, features):
"""Add ops for evaluation (aka filtering) to the graph."""
mode = estimator_lib.ModeKeys.EVAL
with variable_scope.variable_scope("model", use_resource=True):
model_outputs = self.create_loss(features, mode)
metrics = {}
# Just output in-sample predictions for the last chunk seen
for prediction_key, prediction_value in model_outputs.predictions.items(
):
metrics[prediction_key] = _identity_metric_single(
prediction_key, prediction_value)
metrics[feature_keys.FilteringResults.TIMES] = _identity_metric_single(
feature_keys.FilteringResults.TIMES,
model_outputs.prediction_times)
metrics[feature_keys.FilteringResults.STATE_TUPLE] = (
_identity_metric_nested(feature_keys.FilteringResults.STATE_TUPLE,
model_outputs.end_state))
metrics[metric_keys.MetricKeys.LOSS_MEAN] = metrics_impl.mean(
model_outputs.loss, name="average_loss")
return estimator_lib.EstimatorSpec(
loss=model_outputs.loss,
mode=mode,
eval_metric_ops=metrics,
# needed for custom metrics.
predictions=model_outputs.predictions)
def _predict_ops(self, features):
"""Add ops for prediction to the graph."""
with variable_scope.variable_scope("model", use_resource=True):
prediction = self.model.predict(features=features)
prediction[feature_keys.PredictionResults.TIMES] = features[
feature_keys.PredictionFeatures.TIMES]
return estimator_lib.EstimatorSpec(predictions=prediction,
mode=estimator_lib.ModeKeys.PREDICT)
def model_fn(features, labels, mode, params):
"""The model_fn argument for creating an Estimator."""
model = Model(params["data_format"])
image = features
if isinstance(image, dict):
image = features["image"]
if mode == estimator.ModeKeys.PREDICT:
logits = model(image, training=False)
predictions = {
"classes": math_ops.argmax(logits, axis=1),
"probabilities": nn.softmax(logits),
}
return estimator.EstimatorSpec(
mode=estimator.ModeKeys.PREDICT,
predictions=predictions,
export_outputs={
"classify": estimator.export.PredictOutput(predictions)
})
elif mode == estimator.ModeKeys.TRAIN:
optimizer = train.AdamOptimizer(learning_rate=1e-4)
logits = model(image, training=True)
loss = losses.sparse_softmax_cross_entropy(labels=labels,
logits=logits)
return estimator.EstimatorSpec(mode=estimator.ModeKeys.TRAIN,
loss=loss,
train_op=optimizer.minimize(
loss,
train.get_or_create_global_step()))
elif mode == estimator.ModeKeys.EVAL:
logits = model(image, training=False)
loss = losses.sparse_softmax_cross_entropy(labels=labels,
logits=logits)
return estimator.EstimatorSpec(
mode=estimator.ModeKeys.EVAL,
loss=loss,
eval_metric_ops={
"accuracy":
ops.metrics.accuracy(labels=labels,
predictions=math_ops.argmax(logits,
axis=1)),
})
def _train_ops(self, features):
"""Add training ops to the graph."""
mode = estimator_lib.ModeKeys.TRAIN
with variable_scope.variable_scope(
"model",
# Use ResourceVariables to avoid race conditions.
use_resource=True):
model_outputs = self.create_loss(features, mode)
train_op = self.optimizer.minimize(
model_outputs.loss, global_step=training_util.get_global_step())
return estimator_lib.EstimatorSpec(loss=model_outputs.loss,
mode=mode,
train_op=train_op)