(self, device)
| 37 | return create_model(self.cfg.model, class_num=self.class_num, weight_path=self.cfg.weight).to(device) |
| 38 | |
| 39 | def _load_onnx_model(self, device): |
| 40 | from onnxruntime import InferenceSession |
| 41 | |
| 42 | def onnx_forward(self: InferenceSession, x: Tensor): |
| 43 | x = {self.get_inputs()[0].name: x.cpu().numpy()} |
| 44 | model_outputs, layer_output = [], [] |
| 45 | for idx, predict in enumerate(self.run(None, x)): |
| 46 | layer_output.append(torch.from_numpy(predict).to(device)) |
| 47 | if idx % 3 == 2: |
| 48 | model_outputs.append(layer_output) |
| 49 | layer_output = [] |
| 50 | if len(model_outputs) == 6: |
| 51 | model_outputs = model_outputs[:3] |
| 52 | return {"Main": model_outputs} |
| 53 | |
| 54 | InferenceSession.__call__ = onnx_forward |
| 55 | |
| 56 | if device == "cpu": |
| 57 | providers = ["CPUExecutionProvider"] |
| 58 | else: |
| 59 | providers = ["CUDAExecutionProvider"] |
| 60 | try: |
| 61 | ort_session = InferenceSession(self.model_path, providers=providers) |
| 62 | logger.info(":rocket: Using ONNX as MODEL frameworks!") |
| 63 | except Exception as e: |
| 64 | logger.warning(f"🈳 Error loading ONNX model: {e}") |
| 65 | ort_session = self._create_onnx_model(providers) |
| 66 | return ort_session |
| 67 | |
| 68 | def _create_onnx_model(self, providers): |
| 69 | from onnxruntime import InferenceSession |
no test coverage detected