def model_fn(features, labels, mode):
   _, _ = features, labels
   return model_fn_lib.EstimatorSpec(
       mode,
       loss=constant_op.constant([103]),
       train_op=state_ops.assign_add(training.get_global_step(), 1),
       predictions=constant_op.constant([502]),
       export_outputs={'test': export_lib.ClassificationOutput(
           constant_op.constant([[32.]]))})
    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 test_classification_output_tensors_roundtrip_scores_only(self):
        scores = array_ops.placeholder(dtypes.float32, 1, name='scores')
        classification_output = export_lib.ClassificationOutput(scores=scores)

        classification_output_tensors = (
            tpu_estimator._export_output_to_tensors(classification_output))
        self.assertEqual(classification_output_tensors, [scores, None])

        scores_new = array_ops.placeholder(dtypes.float32,
                                           1,
                                           name='scores_new')
        classification_output_new = (
            tpu_estimator._clone_export_output_with_tensors(
                classification_output, [scores_new, None]))
        self.assertEqual(classification_output_new.scores, scores_new)
    def test_classification_output_tensors_roundtrip_classes_only(self):
        classes = array_ops.placeholder(dtypes.string, 1, name='classes')
        classification_output = export_lib.ClassificationOutput(
            classes=classes)

        classification_output_tensors = (
            tpu_estimator._export_output_to_tensors(classification_output))
        self.assertEqual(classification_output_tensors, [None, classes])

        classes_new = array_ops.placeholder(dtypes.string,
                                            1,
                                            name='classes_new')
        classification_output_new = (
            tpu_estimator._clone_export_output_with_tensors(
                classification_output, [None, classes_new]))
        self.assertEqual(classification_output_new.classes, classes_new)
Exemplo n.º 5
0
  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)
Exemplo n.º 6
0
  def test_classification_output_tensors_roundtrip_classify_both(self):
    classes = array_ops.placeholder(dtypes.string, 1, name='classes')
    scores = array_ops.placeholder(dtypes.float32, 1, name='scores')
    classification_output = export_lib.ClassificationOutput(
        scores, classes)

    classification_output_tensors = (tpu_estimator.
                                     _export_output_to_tensors(
                                         classification_output))
    self.assertSequenceEqual(classification_output_tensors, [scores, classes])

    classes_new = array_ops.placeholder(dtypes.string, 1, name='classes_new')
    scores_new = array_ops.placeholder(dtypes.float32, 1, name='scores_new')
    classification_output_new = (tpu_estimator.
                                 _clone_export_output_with_tensors(
                                     classification_output,
                                     [scores_new, classes_new]))
    self.assertEqual(classification_output_new.classes, classes_new)
    self.assertEqual(classification_output_new.scores, scores_new)