(self, providers)
| 66 | return ort_session |
| 67 | |
| 68 | def _create_onnx_model(self, providers): |
| 69 | from onnxruntime import InferenceSession |
| 70 | from torch.onnx import export |
| 71 | |
| 72 | model = create_model(self.cfg.model, class_num=self.class_num, weight_path=self.cfg.weight).eval() |
| 73 | dummy_input = torch.ones((1, 3, *self.cfg.image_size)) |
| 74 | export( |
| 75 | model, |
| 76 | dummy_input, |
| 77 | self.model_path, |
| 78 | input_names=["input"], |
| 79 | output_names=["output"], |
| 80 | dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}, |
| 81 | ) |
| 82 | logger.info(f":inbox_tray: ONNX model saved to {self.model_path}") |
| 83 | return InferenceSession(self.model_path, providers=providers) |
| 84 | |
| 85 | def _load_trt_model(self): |
| 86 | from torch2trt import TRTModule |
no test coverage detected