parameters: (
        self, features, mode, logits, labels=None, train_op_fn=None)
return_statement: return model_fn.EstimatorSpec(
          mode=model_fn.ModeKeys.TRAIN,
          predictions=predictions,
          loss=training_loss,
          train_op=train_op_fn(training_loss))
docstring: See `Head`.
docstring_summary: See `Head`.
docstring_tokens: See Head .
function: 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))
function_tokens: def create_estimator_spec ( self , features , mode , logits , labels = None , train_op_fn = None ) : 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 ) )