(self)
| 88 | self.assertAllClose(output['outputs'].numpy(), expected_output.numpy()) |
| 89 | |
| 90 | def test_processors(self): |
| 91 | model = tf.Module() |
| 92 | inputs = tf.zeros((), tf.float32) |
| 93 | |
| 94 | def _inference_step(inputs, model): |
| 95 | del model |
| 96 | return inputs + 1.0 |
| 97 | |
| 98 | def _preprocessor(inputs): |
| 99 | print(inputs) |
| 100 | return inputs + 0.1 |
| 101 | |
| 102 | module = TestModule( |
| 103 | params=None, |
| 104 | model=model, |
| 105 | inference_step=_inference_step, |
| 106 | preprocessor=_preprocessor) |
| 107 | output = module.serve(inputs) |
| 108 | self.assertAllClose(output['outputs'].numpy(), 1.1) |
| 109 | |
| 110 | class _PostProcessor(tf.Module): |
| 111 | |
| 112 | def __call__(self, inputs): |
| 113 | return inputs + 0.01 |
| 114 | |
| 115 | module = TestModule( |
| 116 | params=None, |
| 117 | model=model, |
| 118 | inference_step=_inference_step, |
| 119 | preprocessor=_preprocessor, |
| 120 | postprocessor=_PostProcessor()) |
| 121 | output = module.serve(inputs) |
| 122 | self.assertAllClose(output['outputs'].numpy(), 1.11) |
| 123 | |
| 124 | def test_get_timestamped_export_dir(self): |
| 125 | export_dir = self.get_temp_dir() |
nothing calls this directly
no test coverage detected