def create_estimator_spec(self, features, mode, logits, labels=None, train_op_fn=None): """See `Head`.""" with ops.name_scope('head'): logits = head_lib._check_logits(logits, self.logits_dimension) # pylint:disable=protected-access # Predict. pred_keys = prediction_keys.PredictionKeys with ops.name_scope(None, 'predictions', (logits, )): probabilities = math_ops.sigmoid(logits, name=pred_keys.PROBABILITIES) predictions = { pred_keys.LOGITS: logits, pred_keys.PROBABILITIES: probabilities, } if mode == model_fn.ModeKeys.PREDICT: return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.PREDICT, predictions=predictions, export_outputs={ '': export_output.ClassificationOutput( scores=probabilities) }) # Eval. unweighted_loss, processed_labels = self.create_loss( features=features, mode=mode, logits=logits, labels=labels) # Averages loss over classes. per_example_loss = math_ops.reduce_mean(unweighted_loss, axis=-1, keep_dims=True) weights = head_lib._weights(features, self._weight_column) # pylint:disable=protected-access training_loss = losses.compute_weighted_loss( per_example_loss, weights=weights, reduction=losses.Reduction.SUM) if mode == model_fn.ModeKeys.EVAL: return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.EVAL, predictions=predictions, loss=training_loss, eval_metric_ops=self._eval_metric_ops( labels=processed_labels, probabilities=probabilities, weights=weights, per_example_loss=per_example_loss)) # Train. if train_op_fn is None: raise ValueError('train_op_fn can not be None.') with ops.name_scope(''): summary.scalar( head_lib._summary_key(self._head_name, metric_keys.MetricKeys.LOSS), # pylint:disable=protected-access training_loss) summary.scalar( head_lib._summary_key( # pylint:disable=protected-access self._head_name, metric_keys.MetricKeys.LOSS_MEAN), losses.compute_weighted_loss(unweighted_loss, weights=weights, reduction=losses.Reduction.MEAN)) return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, predictions=predictions, loss=training_loss, train_op=train_op_fn(training_loss))
def testTrainOpMissing(self): with ops.Graph().as_default(), self.test_session(): with self.assertRaisesRegexp(ValueError, 'Missing train_op'): model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, loss=constant_op.constant(1.))
def testLossMissing(self): with ops.Graph().as_default(), self.test_session(): with self.assertRaisesRegexp(ValueError, 'Missing loss'): model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, train_op=control_flow_ops.no_op())
def create_estimator_spec(self, features, mode, logits, labels=None, train_op_fn=None): """See `Head`.""" with ops.name_scope(self._name, 'head'): logits = head_lib._check_logits(logits, self.logits_dimension) # pylint:disable=protected-access # Predict. pred_keys = prediction_keys.PredictionKeys with ops.name_scope(None, 'predictions', (logits, )): probabilities = math_ops.sigmoid(logits, name=pred_keys.PROBABILITIES) predictions = { pred_keys.LOGITS: logits, pred_keys.PROBABILITIES: probabilities, } if mode == model_fn.ModeKeys.PREDICT: classifier_output = head_lib._classification_output( # pylint:disable=protected-access scores=probabilities, n_classes=self._n_classes, label_vocabulary=self._label_vocabulary) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.PREDICT, predictions=predictions, export_outputs={ _DEFAULT_SERVING_KEY: classifier_output, head_lib._CLASSIFY_SERVING_KEY: classifier_output, # pylint:disable=protected-access head_lib._PREDICT_SERVING_KEY: ( # pylint:disable=protected-access export_output.PredictOutput(predictions)) }) # Eval. unweighted_loss, processed_labels = self.create_loss( features=features, mode=mode, logits=logits, labels=labels) weights = head_lib._weights(features, self._weight_column) # pylint:disable=protected-access training_loss = losses.compute_weighted_loss( unweighted_loss, weights=weights, reduction=losses.Reduction.SUM) if mode == model_fn.ModeKeys.EVAL: return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.EVAL, predictions=predictions, loss=training_loss, eval_metric_ops=self._eval_metric_ops( labels=processed_labels, probabilities=probabilities, weights=weights, unweighted_loss=unweighted_loss)) # Train. if train_op_fn is None: raise ValueError('train_op_fn can not be None.') with ops.name_scope(''): summary.scalar( head_lib._summary_key(self._name, metric_keys.MetricKeys.LOSS), # pylint:disable=protected-access training_loss) summary.scalar( head_lib._summary_key( # pylint:disable=protected-access self._name, metric_keys.MetricKeys.LOSS_MEAN), losses.compute_weighted_loss(unweighted_loss, weights=weights, reduction=losses.Reduction.MEAN)) return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, predictions=predictions, loss=training_loss, train_op=train_op_fn(training_loss))
def testPredictionsTensor(self): """Tests that no error is raised when predictions is Tensor (not dict).""" with ops.Graph().as_default(), self.test_session(): model_fn.EstimatorSpec(mode=model_fn.ModeKeys.PREDICT, predictions=constant_op.constant(1.))
def testRequiredArgumentsSet(self): """Tests that no errors are raised when all required arguments are set.""" with ops.Graph().as_default(), self.test_session(): model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, loss=constant_op.constant(1.), train_op=control_flow_ops.no_op())
def create_estimator_spec( self, features, mode, logits, labels=None, train_op_fn=None, regularization_losses=None): """Returns an `EstimatorSpec`. Args: features: Input `dict` of `Tensor` or `SparseTensor` objects. mode: Estimator's `ModeKeys`. logits: logits `Tensor` with shape `[D0, D1, ... DN, n_classes]`. For many applications, the shape is `[batch_size, n_classes]`. labels: Labels with shape matching `logits`. Can be multi-hot `Tensor` with shape `[D0, D1, ... DN, n_classes]` or `SparseTensor` with `dense_shape` `[D0, D1, ... DN, ?]`. `labels` is required argument when `mode` equals `TRAIN` or `EVAL`. train_op_fn: Function that takes a scalar loss `Tensor` and returns `train_op`. Required in TRAIN mode. regularization_losses: A list of additional scalar losses to be added to the training loss, such as regularization losses. These losses are usually expressed as a batch average, so for best results users need to set `loss_reduction=SUM_OVER_BATCH_SIZE` or `loss_reduction=SUM_OVER_NONZERO_WEIGHTS` when creating the head to avoid scaling errors. Returns: `EstimatorSpec`. Raises: ValueError: If `train_op_fn` is `None` in TRAIN mode. """ with ops.name_scope(self._name, 'head'): logits = head_lib._check_logits_final_dim(logits, self.logits_dimension) # pylint:disable=protected-access # Predict. pred_keys = prediction_keys.PredictionKeys with ops.name_scope(None, 'predictions', (logits,)): probabilities = math_ops.sigmoid(logits, name=pred_keys.PROBABILITIES) predictions = { pred_keys.LOGITS: logits, pred_keys.PROBABILITIES: probabilities, } if mode == model_fn.ModeKeys.PREDICT: classifier_output = head_lib._classification_output( # pylint:disable=protected-access scores=probabilities, n_classes=self._n_classes, label_vocabulary=self._label_vocabulary) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.PREDICT, predictions=predictions, export_outputs={ _DEFAULT_SERVING_KEY: classifier_output, head_lib._CLASSIFY_SERVING_KEY: classifier_output, # pylint:disable=protected-access head_lib._PREDICT_SERVING_KEY: ( # pylint:disable=protected-access export_output.PredictOutput(predictions)) }) (training_loss, unreduced_loss, weights, processed_labels) = self.create_loss( features=features, mode=mode, logits=logits, labels=labels) if regularization_losses: regularization_loss = math_ops.add_n(regularization_losses) regularized_training_loss = math_ops.add_n( [training_loss, regularization_loss]) else: regularization_loss = None regularized_training_loss = training_loss # Eval. if mode == model_fn.ModeKeys.EVAL: return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.EVAL, predictions=predictions, loss=regularized_training_loss, eval_metric_ops=self._eval_metric_ops( labels=processed_labels, probabilities=probabilities, weights=weights, unreduced_loss=unreduced_loss, regularization_loss=regularization_loss)) # Train. if train_op_fn is None: raise ValueError('train_op_fn can not be None.') # Only summarize mean_loss for SUM reduction to preserve backwards # compatibility. Otherwise skip it to avoid unnecessary computation. if self._loss_reduction == losses.Reduction.SUM: example_weight_sum = math_ops.reduce_sum( weights * array_ops.ones_like(unreduced_loss)) mean_loss = training_loss / example_weight_sum else: mean_loss = None with ops.name_scope(''): keys = metric_keys.MetricKeys summary.scalar( head_lib._summary_key(self._name, keys.LOSS), # pylint:disable=protected-access regularized_training_loss) if mean_loss is not None: summary.scalar( head_lib._summary_key(self._name, keys.LOSS_MEAN), # pylint:disable=protected-access mean_loss) if regularization_loss is not None: summary.scalar( head_lib._summary_key(self._name, keys.LOSS_REGULARIZATION), # pylint:disable=protected-access regularization_loss) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.TRAIN, predictions=predictions, loss=regularized_training_loss, train_op=train_op_fn(regularized_training_loss))
def testRequiredArgumentsSet(self): """Tests that no errors are raised when all required arguments are set.""" with ops.Graph().as_default(), self.test_session(): model_fn.EstimatorSpec( mode=model_fn.ModeKeys.PREDICT, predictions={'loss': constant_op.constant(1.)})
def _model_fn(features, labels, mode, params, model): """Model defination for the SSD model based on ResNet-50. Args: features: the input image tensor with shape [batch_size, height, width, 3]. The height and width are fixed and equal. labels: the input labels in a dictionary. The labels include class targets and box targets which are dense label maps. The labels are generated from get_input_fn function in data/dataloader.py mode: the mode of TPUEstimator including TRAIN, EVAL, and PREDICT. params: the dictionary defines hyperparameters of model. The default settings are in default_hparams function in this file. model: the SSD model outputs class logits and box regression outputs. Returns: spec: the EstimatorSpec or TPUEstimatorSpec to run training, evaluation, or prediction. """ if mode == tf.estimator.ModeKeys.PREDICT: labels = features features = labels.pop('image') features -= tf.constant(constants.NORMALIZATION_MEAN, shape=[1, 1, 3], dtype=features.dtype) COEF_STD = 1.0 / tf.constant( constants.NORMALIZATION_STD, shape=[1, 1, 3], dtype=features.dtype) features *= COEF_STD def _model_outputs(): return model(features, params, is_training_bn=(mode == tf.estimator.ModeKeys.TRAIN)) if params['dtype'] == 'bf16': with tf.compat.v1.tpu.bfloat16_scope(): cls_outputs, box_outputs = _model_outputs() levels = cls_outputs.keys() for level in levels: cls_outputs[level] = tf.cast(cls_outputs[level], tf.float32) box_outputs[level] = tf.cast(box_outputs[level], tf.float32) else: cls_outputs, box_outputs = _model_outputs() levels = cls_outputs.keys() # First check if it is in PREDICT mode. if mode == tf.estimator.ModeKeys.PREDICT: flattened_cls, flattened_box = concat_outputs(cls_outputs, box_outputs, True) ssd_box_coder = faster_rcnn_box_coder.FasterRcnnBoxCoder( scale_factors=constants.BOX_CODER_SCALES) anchors = box_list.BoxList( tf.convert_to_tensor(dataloader.DefaultBoxes()('ltrb'))) decoded_boxes = box_coder.batch_decode(encoded_boxes=flattened_box, box_coder=ssd_box_coder, anchors=anchors) pred_scores = tf.nn.softmax(flattened_cls, axis=2) pred_scores, indices = select_top_k_scores( pred_scores, constants.MAX_NUM_EVAL_BOXES) predictions = dict( labels, indices=indices, pred_scores=pred_scores, pred_box=decoded_boxes, ) if params['visualize_dataloader']: # this is for inference visualization. predictions['image'] = features return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions) # Load pretrained model from checkpoint. if params['resnet_checkpoint'] and mode == tf.estimator.ModeKeys.TRAIN: def scaffold_fn(): """Loads pretrained model through scaffold function.""" tf.train.init_from_checkpoint( params['resnet_checkpoint'], { '/': 'resnet%s/' % constants.RESNET_DEPTH, }) return tf.train.Scaffold() else: scaffold_fn = None # Set up training loss and learning rate. update_learning_rate_schedule_parameters(params) global_step = tf.train.get_or_create_global_step() learning_rate = learning_rate_schedule(params, global_step) # cls_loss and box_loss are for logging. only total_loss is optimized. loss, cls_loss, box_loss = detection_loss(cls_outputs, box_outputs, labels) total_loss = loss + params['weight_decay'] * tf.add_n( [tf.nn.l2_loss(v) for v in tf.trainable_variables()]) if mode == tf.estimator.ModeKeys.TRAIN: optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=constants.MOMENTUM) if params['distributed_optimizer']: optimizer = params['distributed_optimizer'](optimizer) # Batch norm requires update_ops to be added as a train_op dependency. update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) train_op = tf.group(optimizer.minimize(total_loss, global_step), update_ops) save_summary_steps = params['save_summary_steps'] if save_summary_steps is not None: tf.summary.scalar("learning_rate", learning_rate) tf.summary.scalar("class_loss", cls_loss) tf.summary.scalar("box_loss", box_loss) return model_fn_lib.EstimatorSpec(mode=mode, loss=loss, train_op=train_op, training_hooks=None, scaffold=scaffold_fn()) if mode == tf.estimator.ModeKeys.EVAL: raise NotImplementedError
def estimator_spec(self, mode, default_serving_output_alternative_key=None): """Creates an equivalent `EstimatorSpec`. Args: mode: One of `ModeKeys`. Specifies if this training, evaluation or prediction. default_serving_output_alternative_key: Required for multiple heads. If you have multiple entries in `output_alternatives` dict (comparable to multiple heads), `EstimatorSpec` requires a default head that will be used if a Servo request does not explicitly mention which head to infer on. Pass the key of the output alternative here that you want to designate as default. A separate ExportOutpout for this default head wil be added to the export_outputs dict with the special key signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, unless there is already an enry in output_alternatives with this special key. Returns: Instance of `EstimatorSpec` that is equivalent to this `ModelFnOps` Raises: ValueError: If problem type is unknown. """ def _scores(output_tensors): scores = output_tensors.get(prediction_key.PredictionKey.SCORES) if scores is None: scores = output_tensors.get( prediction_key.PredictionKey.PROBABILITIES) return scores def _classes(output_tensors): # pylint: disable=missing-docstring classes = output_tensors.get(prediction_key.PredictionKey.CLASSES) if classes is None: logging.warning( 'classes is None, Servo inference will not have class ids.' ) return None elif classes.dtype != dtypes.string: # Servo classification can only serve string classes logging.warning( 'classes is not string, Servo inference will not have class ids.' ) return None return classes def _export_output(problem_type, predictions): # pylint: disable=missing-docstring if problem_type == constants.ProblemType.LINEAR_REGRESSION: return core_export_lib.RegressionOutput(_scores(predictions)) if (problem_type == constants.ProblemType.CLASSIFICATION or problem_type == constants.ProblemType.LOGISTIC_REGRESSION): return core_export_lib.ClassificationOutput( scores=_scores(predictions), classes=_classes(predictions)) if problem_type == constants.ProblemType.UNSPECIFIED: return core_export_lib.PredictOutput(predictions) raise ValueError('Unknown problem_type=%s' % problem_type) # Converts output_alternatives export_outputs_dict = None if self.output_alternatives: output_alternatives = self.output_alternatives # Adds default output_alternative if needed. if (len(output_alternatives) > 1 and signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY not in output_alternatives): output_alternatives = output_alternatives.copy() output_alternatives[ signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = ( output_alternatives[ default_serving_output_alternative_key]) export_outputs_dict = { key: _export_output(*val) for key, val in output_alternatives.items() } def _get_eval_metric_ops(): """Returns self.eval_metric_ops without loss metric.""" result = {} for key, value in six.iteritems(self.eval_metric_ops): if key != metric_key.MetricKey.LOSS: result[key] = value return result return core_model_fn_lib.EstimatorSpec( mode=mode, predictions=self.predictions, loss=self.loss, train_op=self.train_op, eval_metric_ops=_get_eval_metric_ops(), export_outputs=export_outputs_dict, training_chief_hooks=self.training_chief_hooks, training_hooks=self.training_hooks, scaffold=self.scaffold)
def my_model_fn(features, mode): loss = math_ops.reduce_sum(features) train_op = array_ops.identity(loss) return model_fn_lib.EstimatorSpec(mode, loss=loss, train_op=train_op)
def create_estimator_spec(self, features, mode, logits, labels=None, train_op_fn=None): """See `Head`.""" # Predict. with ops.name_scope(self._name, 'head'): logits = _check_logits(logits, self._logits_dimension) predictions = {prediction_keys.PredictionKeys.PREDICTIONS: logits} if mode == model_fn.ModeKeys.PREDICT: regression_output = export_output.RegressionOutput( value=logits) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.PREDICT, predictions=predictions, export_outputs={ _DEFAULT_SERVING_KEY: regression_output, _REGRESS_SERVING_KEY: regression_output, _PREDICT_SERVING_KEY: export_output.PredictOutput(predictions) }) # Eval. unweighted_loss, _ = self.create_loss(features=features, mode=mode, logits=logits, labels=labels) weights = _weights(features, self._weight_column) training_loss = losses.compute_weighted_loss( unweighted_loss, weights=weights, reduction=losses.Reduction.SUM) if mode == model_fn.ModeKeys.EVAL: # Estimator already adds a metric for loss. eval_metric_ops = { metric_keys.MetricKeys.LOSS_MEAN: metrics_lib.mean(unweighted_loss, weights=weights) } return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.EVAL, predictions=predictions, loss=training_loss, eval_metric_ops=eval_metric_ops) # Train. if train_op_fn is None: raise ValueError('train_op_fn can not be None.') with ops.name_scope(''): summary.scalar( _summary_key(self._name, metric_keys.MetricKeys.LOSS), training_loss) summary.scalar( _summary_key(self._name, metric_keys.MetricKeys.LOSS_MEAN), losses.compute_weighted_loss(unweighted_loss, weights=weights, reduction=losses.Reduction.MEAN)) return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, predictions=predictions, loss=training_loss, train_op=train_op_fn(training_loss))
def create_estimator_spec(self, features, mode, logits, labels=None, train_op_fn=None): """See `Head`.""" # Predict. with ops.name_scope(self._name, 'head'): with ops.name_scope(None, 'predictions', (logits, )): pred_keys = prediction_keys.PredictionKeys logits = _check_logits(logits, self.logits_dimension) logistic = math_ops.sigmoid(logits, name=pred_keys.LOGISTIC) two_class_logits = array_ops.concat( (array_ops.zeros_like(logits), logits), 1, name='two_class_logits') probabilities = nn.softmax(two_class_logits, name=pred_keys.PROBABILITIES) class_ids = array_ops.reshape(math_ops.argmax(two_class_logits, axis=1), (-1, 1), name='classes') if self._label_vocabulary: table = lookup_ops.index_to_string_table_from_tensor( vocabulary_list=self._label_vocabulary, name='class_string_lookup') classes = table.lookup(class_ids) else: classes = string_ops.as_string(class_ids, name='str_classes') predictions = { pred_keys.LOGITS: logits, pred_keys.LOGISTIC: logistic, pred_keys.PROBABILITIES: probabilities, pred_keys.CLASS_IDS: class_ids, pred_keys.CLASSES: classes, } if mode == model_fn.ModeKeys.PREDICT: classifier_output = _classification_output( scores=probabilities, n_classes=2, label_vocabulary=self._label_vocabulary) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.PREDICT, predictions=predictions, export_outputs={ _DEFAULT_SERVING_KEY: classifier_output, _CLASSIFY_SERVING_KEY: classifier_output, _REGRESS_SERVING_KEY: export_output.RegressionOutput(value=logistic), _PREDICT_SERVING_KEY: export_output.PredictOutput(predictions) }) # Eval. unweighted_loss, processed_labels = self.create_loss( features=features, mode=mode, logits=logits, labels=labels) weights = _weights(features, self._weight_column) training_loss = losses.compute_weighted_loss( unweighted_loss, weights=weights, reduction=losses.Reduction.SUM) if mode == model_fn.ModeKeys.EVAL: return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.EVAL, predictions=predictions, loss=training_loss, eval_metric_ops=self._eval_metric_ops( labels=processed_labels, logits=logits, logistic=logistic, class_ids=class_ids, unweighted_loss=unweighted_loss, weights=weights)) # Train. if train_op_fn is None: raise ValueError('train_op_fn can not be None.') with ops.name_scope(''): summary.scalar( _summary_key(self._name, metric_keys.MetricKeys.LOSS), training_loss) summary.scalar( _summary_key(self._name, metric_keys.MetricKeys.LOSS_MEAN), losses.compute_weighted_loss(unweighted_loss, weights=weights, reduction=losses.Reduction.MEAN)) return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, predictions=predictions, loss=training_loss, train_op=train_op_fn(training_loss))
def testLossMissing(self): with ops.Graph().as_default(), self.test_session(): with self.assertRaisesRegexp(ValueError, 'Missing loss'): model_fn.EstimatorSpec( mode=model_fn.ModeKeys.EVAL, predictions={'loss': constant_op.constant(1.)})
def model_fn(features, labels, mode): """model_fn for keras Estimator.""" 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 distribute_lib.has_distribution_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 model_fn_lib.ModeKeys.PREDICT: if mode is model_fn_lib.ModeKeys.TRAIN: model._make_train_function() # pylint: disable=protected-access else: model._make_test_function() # pylint: disable=protected-access loss = model.total_loss if model.metrics: # TODO(fchollet): support stateful metrics eval_metric_ops = {} # When each metric maps to an output if isinstance(model.metrics, dict): for i, output_name in enumerate(model.metrics.keys()): metric_name = model.metrics[output_name] if callable(metric_name): metric_name = metric_name.__name__ # When some outputs use the same metric if list(model.metrics.values()).count(metric_name) > 1: metric_name += '_' + output_name eval_metric_ops[metric_name] = metrics_module.mean( model.metrics_tensors[i - len(model.metrics)]) else: for i, metric_name in enumerate(model.metrics): if callable(metric_name): metric_name = metric_name.__name__ eval_metric_ops[metric_name] = metrics_module.mean( model.metrics_tensors[i]) # Set train_op only during train. if mode is model_fn_lib.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. _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.export_output.PredictOutput(predictions) })
def testPredictionsMissingIsOkay(self): with ops.Graph().as_default(), self.test_session(): model_fn.EstimatorSpec(mode=model_fn.ModeKeys.EVAL, loss=constant_op.constant(1.))
def eval_model_fn_no_eval_metric_ops(features, labels, mode, params): del features, labels, params return model_fn_lib.EstimatorSpec( mode=mode, loss=constant_op.constant(_EXPECTED_LOSS))
def testPredictionsMissing(self): with ops.Graph().as_default(), self.test_session(): with self.assertRaisesRegexp(ValueError, 'Missing predictions'): model_fn.EstimatorSpec(mode=model_fn.ModeKeys.PREDICT)
def _test_train_model_fn(features, labels, mode, params): """A dummy model_fn for testing purpose.""" del features, labels, params loss = constant_op.constant(_EXPECTED_LOSS) return model_fn_lib.EstimatorSpec( mode=mode, loss=loss, train_op=array_ops.identity(loss))
def testPredictionsNumber(self): with ops.Graph().as_default(), self.test_session(): with self.assertRaisesRegexp( TypeError, r'predictions\[number\] must be Tensor'): model_fn.EstimatorSpec(mode=model_fn.ModeKeys.PREDICT, predictions={'number': 1.})
def create_estimator_spec(self, features, mode, logits, labels=None, train_op_fn=None): """See `Head`.""" with ops.name_scope('head'): logits = _check_logits(logits, self.logits_dimension) # Predict. pred_keys = prediction_keys.PredictionKeys with ops.name_scope(None, 'predictions', (logits, )): # class_ids's shape is [batch_size] class_ids = math_ops.argmax(logits, 1, name=pred_keys.CLASS_IDS) class_ids = array_ops.expand_dims(class_ids, axis=(1, )) if self._label_vocabulary: table = lookup_ops.index_to_string_table_from_tensor( vocabulary_list=self._label_vocabulary, name='class_string_lookup') classes = table.lookup(class_ids) else: classes = string_ops.as_string(class_ids, name='str_classes') probabilities = nn.softmax(logits, name=pred_keys.PROBABILITIES) predictions = { pred_keys.LOGITS: logits, pred_keys.PROBABILITIES: probabilities, # Expand to [batch_size, 1] pred_keys.CLASS_IDS: class_ids, pred_keys.CLASSES: classes, } if mode == model_fn.ModeKeys.PREDICT: batch_size = array_ops.shape(probabilities)[0] export_class_list = self._label_vocabulary if not export_class_list: export_class_list = string_ops.as_string( math_ops.range(self._n_classes)) export_output_classes = array_ops.tile( input=array_ops.expand_dims(input=export_class_list, axis=0), multiples=[batch_size, 1]) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.PREDICT, predictions=predictions, export_outputs={ '': export_output.ClassificationOutput( scores=probabilities, # `ClassificationOutput` requires string classes. classes=export_output_classes) }) # Eval. label_ids = self._label_ids( _check_labels(_maybe_expand_dim(labels), 1)) unweighted_loss = losses.sparse_softmax_cross_entropy( labels=label_ids, logits=logits, reduction=losses.Reduction.NONE) # Restore the squeezed dim, so unweighted_loss matches the weights shape. unweighted_loss = array_ops.expand_dims(unweighted_loss, axis=(1, )) weights = _weights(features, self._weight_column) training_loss = losses.compute_weighted_loss( unweighted_loss, weights=weights, reduction=losses.Reduction.SUM) if mode == model_fn.ModeKeys.EVAL: return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.EVAL, predictions=predictions, loss=training_loss, eval_metric_ops=self._eval_metric_ops( labels=label_ids, probabilities=probabilities, logits=logits, class_ids=class_ids, unweighted_loss=unweighted_loss, weights=weights)) # Train. if train_op_fn is None: raise ValueError('train_op_fn can not be None.') with ops.name_scope(''): summary.scalar(metric_keys.MetricKeys.LOSS, training_loss) summary.scalar( metric_keys.MetricKeys.LOSS_MEAN, losses.compute_weighted_loss(unweighted_loss, weights=weights, reduction=losses.Reduction.MEAN)) return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, predictions=predictions, loss=training_loss, train_op=train_op_fn(training_loss))
def testLoss1DTensor(self): """Tests that no errors are raised when loss is 1D tensor.""" with ops.Graph().as_default(), self.test_session(): model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, loss=constant_op.constant([1.]), train_op=control_flow_ops.no_op())
def create_estimator_spec(self, features, mode, logits, labels=None, train_op_fn=None): """See `Head`.""" # Predict. with ops.name_scope('head'): with ops.name_scope(None, 'predictions', (logits, )): pred_keys = prediction_keys.PredictionKeys logits = _check_logits(logits, self.logits_dimension) logistic = math_ops.sigmoid(logits, name=pred_keys.LOGISTIC) two_class_logits = array_ops.concat( (array_ops.zeros_like(logits), logits), 1, name='two_class_logits') scores = nn.softmax(two_class_logits, name=pred_keys.PROBABILITIES) class_ids = array_ops.reshape(math_ops.argmax(two_class_logits, axis=1), (-1, 1), name='classes') if self._label_vocabulary: table = lookup_ops.index_to_string_table_from_tensor( vocabulary_list=self._label_vocabulary, name='class_string_lookup') classes = table.lookup(class_ids) else: classes = string_ops.as_string(class_ids, name='str_classes') predictions = { pred_keys.LOGITS: logits, pred_keys.LOGISTIC: logistic, pred_keys.PROBABILITIES: scores, pred_keys.CLASS_IDS: class_ids, pred_keys.CLASSES: classes, } if mode == model_fn.ModeKeys.PREDICT: batch_size = array_ops.shape(logistic)[0] export_class_list = self._label_vocabulary if not export_class_list: export_class_list = string_ops.as_string([0, 1]) export_output_classes = array_ops.tile( input=array_ops.expand_dims(input=export_class_list, axis=0), multiples=[batch_size, 1]) classifier_output = export_output.ClassificationOutput( scores=scores, # `ClassificationOutput` requires string classes. classes=export_output_classes) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.PREDICT, predictions=predictions, export_outputs={ '': classifier_output, # to be same as other heads. 'classification': classifier_output, # to be called by name. _DEFAULT_SERVING_KEY: classifier_output, # default 'regression': export_output.RegressionOutput(value=logistic) }) # Eval. labels = _check_labels(_maybe_expand_dim(labels), self.logits_dimension) if self._label_vocabulary is not None: labels = lookup_ops.index_table_from_tensor( vocabulary_list=tuple(self._label_vocabulary), name='class_id_lookup').lookup(labels) labels = math_ops.to_float(labels) labels = _assert_range(labels, 2) unweighted_loss = nn.sigmoid_cross_entropy_with_logits( labels=labels, logits=logits, name='loss') weights = _weights(features, self._weight_column) training_loss = losses.compute_weighted_loss( unweighted_loss, weights=weights, reduction=losses.Reduction.SUM) if mode == model_fn.ModeKeys.EVAL: return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.EVAL, predictions=predictions, loss=training_loss, eval_metric_ops=self._eval_metric_ops( labels=labels, logits=logits, logistic=logistic, scores=scores, class_ids=class_ids, unweighted_loss=unweighted_loss, weights=weights)) # Train. if train_op_fn is None: raise ValueError('train_op_fn can not be None.') with ops.name_scope(''): summary.scalar(metric_keys.MetricKeys.LOSS, training_loss) summary.scalar( metric_keys.MetricKeys.LOSS_MEAN, losses.compute_weighted_loss(unweighted_loss, weights=weights, reduction=losses.Reduction.MEAN)) return model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, predictions=predictions, loss=training_loss, train_op=train_op_fn(training_loss))
def testLossNotScalar(self): with ops.Graph().as_default(), self.test_session(): with self.assertRaisesRegexp(ValueError, 'Loss must be scalar'): model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN, loss=constant_op.constant([1., 2.]), train_op=control_flow_ops.no_op())
def create_estimator_spec(self, features, mode, logits, labels=None, train_op_fn=tfgan_train.gan_train_ops): """Returns `EstimatorSpec` that a model_fn can return. See `Head` for more details. Args: features: Must be `None`. mode: Estimator's `ModeKeys`. logits: A GANModel tuple. labels: Must be `None`. train_op_fn: Function that takes a GANModel, GANLoss, generator optimizer, and discriminator optimizer, and returns a `GANTrainOps` tuple. For example, this function can come from TFGAN's `train.py` library, or can be custom. Returns: `EstimatorSpec`. Raises: ValueError: If `features` isn't `None`. ValueError: If `train_op_fn` isn't provided in train mode. """ _validate_logits_and_labels(logits, labels) if features is not None: raise ValueError( '`features` should be `None`. Instead, found: %s' % features) gan_model = logits # rename variable for clarity with ops.name_scope('GANHead'): if mode == model_fn_lib.ModeKeys.PREDICT: return model_fn_lib.EstimatorSpec( mode=model_fn_lib.ModeKeys.PREDICT, predictions=gan_model.generated_data, export_outputs={ 'predict': export_output.PredictOutput(gan_model.generated_data) }) elif mode == model_fn_lib.ModeKeys.EVAL: gan_loss = self.create_loss(features=None, mode=mode, logits=gan_model, labels=None) scalar_loss = gan_loss.generator_loss + gan_loss.discriminator_loss with ops.name_scope( None, 'metrics', [gan_loss.generator_loss, gan_loss.discriminator_loss]): eval_metric_ops = { _summary_key(self._name, 'generator_loss'): metrics_lib.mean(gan_loss.generator_loss), _summary_key(self._name, 'discriminator_loss'): metrics_lib.mean(gan_loss.discriminator_loss) } if self._get_eval_metric_ops_fn is not None: custom_eval_metric_ops = self._get_eval_metric_ops_fn( gan_model) if not isinstance(custom_eval_metric_ops, dict): raise TypeError( 'get_eval_metric_ops_fn must return a dict, ' 'received: {}'.format(custom_eval_metric_ops)) eval_metric_ops.update(custom_eval_metric_ops) return model_fn_lib.EstimatorSpec( mode=model_fn_lib.ModeKeys.EVAL, predictions=gan_model.generated_data, loss=scalar_loss, eval_metric_ops=eval_metric_ops) elif mode == model_fn_lib.ModeKeys.TRAIN: if train_op_fn is None: raise ValueError('train_op_fn can not be None.') gan_loss = self.create_loss(None, mode, gan_model, None) scalar_loss = gan_loss.generator_loss + gan_loss.discriminator_loss train_ops = train_op_fn(gan_model, gan_loss, self._generator_optimizer, self._discriminator_optimizer) training_hooks = self._get_hooks_fn(train_ops) return model_fn_lib.EstimatorSpec( loss=scalar_loss, mode=model_fn_lib.ModeKeys.TRAIN, train_op=train_ops.global_step_inc_op, training_hooks=training_hooks) else: raise ValueError('Mode not recognized: %s' % mode)