def test_external_streaming_shapes(self, model_name): model = utils.get_model_with_default_params( model_name, mode=modes.Modes.STREAM_EXTERNAL_STATE_INFERENCE) # The first 'n' inputs correspond to the 'n' inputs that the model takes # in non-streaming mode. The rest of the input tensors represent the # internal states for each layer in the model. inputs = [np.zeros(shape, dtype=np.float32) for shape in model.input_shapes] outputs = model.predict(inputs) for output, expected_shape in zip(outputs, model.output_shapes): self.assertEqual(output.shape, expected_shape)
def __init__(self): super().__init__() self.m = utils.get_model_with_default_params( FLAGS.model, MODE_ENUM_TO_MODE[FLAGS.mode]) call = lambda *args: self.m(*args, training=False) input_signature = [ tf.TensorSpec(tensor.shape) for tensor in self.m.inputs ] self.call = tf_test_utils.tf_function_unit_test( input_signature=input_signature, name="call", atol=1e-5)(call)
def create_module(cls): model = utils.get_model_with_default_params( FLAGS.model, MODE_ENUM_TO_MODE[FLAGS.mode]) cls.input_shapes = [tensor.shape for tensor in model.inputs] return cls(model)