def _serving_ops(self, features): """Add ops for serving to the graph.""" with tf.compat.v1.variable_scope("model", use_resource=True): prediction_outputs = self.model.predict(features=features) with tf.compat.v1.variable_scope("model", reuse=True): filtering_outputs = self.create_loss(features, estimator_lib.ModeKeys.EVAL) with tf.compat.v1.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 model_fn(features, labels, mode): """model_fn for keras Estimator.""" # Raise an error when users use DistributionStrategy with native Keras # optimizers. Currently we only support native TensorFlow optimizers. if distribution_strategy_context.has_strategy() and \ not isinstance(keras_model.optimizer, (tf_optimizer_module.Optimizer, optimizers.TFOptimizer)): raise ValueError( 'Only TensorFlow native optimizers are supported with ' 'DistributionStrategy.') model = _clone_and_build_model(mode, keras_model, custom_objects, features, labels) model_output_names = [] # We need to make sure that the output names of the last layer in the model # is the same for each of the cloned models. This is required for mirrored # strategy when we call regroup. if distribution_strategy_context.has_strategy(): for name in model.output_names: name = re.compile(r'_\d$').sub('', name) model_output_names.append(name) else: model_output_names = model.output_names # Get inputs to EstimatorSpec predictions = dict(zip(model_output_names, model.outputs)) loss = None train_op = None eval_metric_ops = None # Set loss and metric only during train and evaluate. if mode is not ModeKeys.PREDICT: if mode is ModeKeys.TRAIN: model._make_train_function() # pylint: disable=protected-access else: model._make_test_function() # pylint: disable=protected-access loss = model.total_loss eval_metric_ops = _convert_keras_metrics_to_estimator(model) # Set train_op only during train. if mode is ModeKeys.TRAIN: train_op = model.train_function.updates_op if not model._is_graph_network: # Reset model state to original state, # to avoid `model_fn` being destructive for the initial model argument. models.in_place_subclassed_model_state_restoration(keras_model) return model_fn_lib.EstimatorSpec( mode=mode, predictions=predictions, loss=loss, train_op=train_op, eval_metric_ops=eval_metric_ops, export_outputs={ _DEFAULT_SERVING_KEY: export_lib.PredictOutput(predictions) })
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, 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_predict_output_tensors_roundtrip(self): value1 = array_ops.placeholder(dtypes.float32, 1, name='value1') value2 = array_ops.placeholder(dtypes.float32, 1, name='value2') predict_output = export_lib.PredictOutput({ 'value1': value1, 'value2': value2 }) export_output_tensors = tpu_estimator._export_output_to_tensors( predict_output) self.assertSameElements([value1, value2], export_output_tensors) self.assertLen(export_output_tensors, 2) tensors_new = [ array_ops.identity(t, name=t.name.split(':')[0] + '_new') for t in export_output_tensors ] predict_output_new = tpu_estimator._clone_export_output_with_tensors( predict_output, tensors_new) outputs = predict_output_new.outputs self.assertLen(outputs, 2) self.assertEqual(outputs['value1'].name, 'value1_new:0') self.assertEqual(outputs['value2'].name, 'value2_new:0')
def model_fn(features, labels, mode): """model_fn for keras Estimator.""" model = _clone_and_build_model(mode=mode, keras_model=keras_model, custom_objects=custom_objects, features=features, labels=labels, optimizer_config=optimizer_config) model_output_names = [] # We need to make sure that the output names of the last layer in the model # is the same for each of the cloned models. This is required for mirrored # strategy when we call regroup. if tf.distribute.has_strategy(): for name in model.output_names: name = re.compile(r'_\d$').sub('', name) model_output_names.append(name) else: model_output_names = model.output_names # Get inputs to EstimatorSpec predictions = dict(zip(model_output_names, model.outputs)) loss = None train_op = None eval_metric_ops = None # Set loss and metric only during train and evaluate. if mode is not ModeKeys.PREDICT: if mode is ModeKeys.TRAIN: model._make_train_function() # pylint: disable=protected-access else: model._make_test_function() # pylint: disable=protected-access loss = model.total_loss eval_metric_ops = _convert_keras_metrics_to_estimator(model) # Set train_op only during train. if mode is ModeKeys.TRAIN: train_op = model.train_function.updates_op if (not model._is_graph_network and hasattr(keras_model, '_original_attributes_cache') and keras_model._original_attributes_cache is not None): # To avoid `model_fn` being destructive for the initial model argument. models.in_place_subclassed_model_state_restoration(keras_model) scaffold = None if save_object_ckpt: model._track_trackable(tf.compat.v1.train.get_global_step(), 'estimator_global_step') # Create saver that maps variable names to object-checkpoint keys. object_graph = graph_view.ObjectGraphView(model) var_list = object_graph.frozen_saveable_objects() saver = tf.compat.v1.train.Saver(var_list=var_list, sharded=True) saver._object_restore_saver = trackable_util.frozen_saver(model) scaffold = tf.compat.v1.train.Scaffold(saver=saver) return model_fn_lib.EstimatorSpec( mode=mode, predictions=predictions, loss=loss, train_op=train_op, eval_metric_ops=eval_metric_ops, export_outputs={ _DEFAULT_SERVING_KEY: export_lib.PredictOutput(predictions) }, scaffold=scaffold)