(self, convert_precision, input_shape, dynamic_batch)
| 70 | @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4]) |
| 71 | @unittest.skipUnless(has_torchtrt and has_tensorrt, "Torch-TensorRT is required for conversion!") |
| 72 | def test_trt_export(self, convert_precision, input_shape, dynamic_batch): |
| 73 | tests_dir = Path(__file__).resolve().parents[1] |
| 74 | meta_file = os.path.join(tests_dir, "testing_data", "metadata.json") |
| 75 | config_file = os.path.join(tests_dir, "testing_data", "inference.json") |
| 76 | with tempfile.TemporaryDirectory() as tempdir: |
| 77 | def_args = {"meta_file": "will be replaced by `meta_file` arg"} |
| 78 | def_args_file = os.path.join(tempdir, "def_args.yaml") |
| 79 | |
| 80 | ckpt_file = os.path.join(tempdir, "model.pt") |
| 81 | ts_file = os.path.join(tempdir, f"model_trt_{convert_precision}.ts") |
| 82 | |
| 83 | parser = ConfigParser() |
| 84 | parser.export_config_file(config=def_args, filepath=def_args_file) |
| 85 | parser.read_config(config_file) |
| 86 | net = parser.get_parsed_content("network_def") |
| 87 | save_state(src=net, path=ckpt_file) |
| 88 | |
| 89 | cmd = ["python", "-m", "monai.bundle", "trt_export", "network_def", "--filepath", ts_file] |
| 90 | cmd += ["--meta_file", meta_file, "--config_file", f"['{config_file}','{def_args_file}']", "--ckpt_file"] |
| 91 | cmd += [ckpt_file, "--args_file", def_args_file, "--precision", convert_precision] |
| 92 | if input_shape: |
| 93 | cmd += ["--input_shape", str(input_shape)] |
| 94 | if dynamic_batch: |
| 95 | cmd += ["--dynamic_batch", str(dynamic_batch)] |
| 96 | command_line_tests(cmd) |
| 97 | self.assertTrue(os.path.exists(ts_file)) |
| 98 | |
| 99 | _, metadata, extra_files = load_net_with_metadata( |
| 100 | ts_file, more_extra_files=["inference.json", "def_args.json"] |
| 101 | ) |
| 102 | self.assertIn("schema", metadata) |
| 103 | self.assertIn("meta_file", json.loads(extra_files["def_args.json"])) |
| 104 | self.assertIn("network_def", json.loads(extra_files["inference.json"])) |
| 105 | |
| 106 | @parameterized.expand([TEST_CASE_3, TEST_CASE_4]) |
| 107 | @unittest.skipUnless( |
nothing calls this directly
no test coverage detected