| 9 | |
| 10 | |
| 11 | class FastModelLoader: |
| 12 | def __init__(self, cfg: Config): |
| 13 | self.cfg = cfg |
| 14 | self.compiler = cfg.task.fast_inference |
| 15 | self.class_num = cfg.dataset.class_num |
| 16 | |
| 17 | self._validate_compiler() |
| 18 | if cfg.weight == True: |
| 19 | cfg.weight = Path("weights") / f"{cfg.model.name}.pt" |
| 20 | self.model_path = f"{Path(cfg.weight).stem}.{self.compiler}" |
| 21 | |
| 22 | def _validate_compiler(self): |
| 23 | if self.compiler not in ["onnx", "trt", "deploy"]: |
| 24 | logger.warning(f":warning: Compiler '{self.compiler}' is not supported. Using original model.") |
| 25 | self.compiler = None |
| 26 | if self.cfg.device == "mps" and self.compiler == "trt": |
| 27 | logger.warning(":red_apple: TensorRT does not support MPS devices. Using original model.") |
| 28 | self.compiler = None |
| 29 | |
| 30 | def load_model(self, device): |
| 31 | if self.compiler == "onnx": |
| 32 | return self._load_onnx_model(device) |
| 33 | elif self.compiler == "trt": |
| 34 | return self._load_trt_model().to(device) |
| 35 | elif self.compiler == "deploy": |
| 36 | self.cfg.model.model.auxiliary = {} |
| 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): |