Class TestModule

official/core/export_base_test.py:24–37

class TestModule(export_base.ExportModule):

  @tf.function
  def serve(self, inputs: tf.Tensor) -> Mapping[Text, tf.Tensor]:
    x = inputs if self.preprocessor is None else self.preprocessor(
        inputs=inputs)
    x = self.inference_step(x)
    x = self.postprocessor(x) if self.postprocessor else x
    return {'outputs': x}

  def get_inference_signatures(
      self, function_keys: Dict[Text, Text]) -> Mapping[Text, Any]:
    input_signature = tf.TensorSpec(shape=[None, None], dtype=tf.float32)
    return {'foo': self.serve.get_concrete_function(input_signature)}


class ExportBaseTest(tf.test.TestCase):