Main function to export the model.
()
| 193 | |
| 194 | @torch.no_grad() |
| 195 | def main(): |
| 196 | """Main function to export the model.""" |
| 197 | parser = get_parser() |
| 198 | args = parser.parse_args() |
| 199 | |
| 200 | tensorrt_model_dir = Path(args.tensorrt_model_dir) |
| 201 | tensorrt_model_dir.mkdir(parents=True, exist_ok=True) |
| 202 | |
| 203 | if args.onnx_model_path: |
| 204 | logging.info(f"Using provided ONNX model: {args.onnx_model_path}") |
| 205 | if not args.idim: |
| 206 | raise ValueError("--idim is required when using --onnx-model-path") |
| 207 | idim = args.idim |
| 208 | encoder_onnx_file = Path(args.onnx_model_path) |
| 209 | if not encoder_onnx_file.is_file(): |
| 210 | raise FileNotFoundError(f"ONNX model not found at {encoder_onnx_file}") |
| 211 | else: |
| 212 | from fireredasr.models.fireredasr import load_fireredasr_aed_model |
| 213 | if not args.model_dir: |
| 214 | raise ValueError( |
| 215 | "--model-dir is required if --onnx-model-path is not provided" |
| 216 | ) |
| 217 | |
| 218 | logging.info("Exporting ONNX model from PyTorch checkpoint") |
| 219 | model_dir = Path(args.model_dir) |
| 220 | model_path = model_dir / "model.pth.tar" |
| 221 | |
| 222 | # Load model to get encoder |
| 223 | package = torch.load(model_path, map_location="cpu", weights_only=False) |
| 224 | model_args = package["args"] |
| 225 | idim = model_args.idim |
| 226 | # We have to load the full AED model to get the encoder with weights |
| 227 | model = load_fireredasr_aed_model(str(model_path)) |
| 228 | encoder = model.encoder |
| 229 | encoder.eval() |
| 230 | |
| 231 | # Export ONNX |
| 232 | encoder_onnx_file = tensorrt_model_dir / "encoder.fp16.onnx" |
| 233 | export_encoder_onnx( |
| 234 | encoder=encoder, |
| 235 | filename=str(encoder_onnx_file), |
| 236 | idim=idim, |
| 237 | opset_version=args.opset_version, |
| 238 | ) |
| 239 | |
| 240 | # Convert ONNX to TensorRT |
| 241 | trt_engine_file = tensorrt_model_dir / args.trt_engine_file_name |
| 242 | trt_kwargs = get_trt_kwargs_dynamic_batch(idim=idim) |
| 243 | convert_onnx_to_trt( |
| 244 | trt_model=str(trt_engine_file), |
| 245 | trt_kwargs=trt_kwargs, |
| 246 | onnx_model=str(encoder_onnx_file), |
| 247 | dtype=torch.float16, |
| 248 | ) |
| 249 | |
| 250 | logging.info("Done!") |
| 251 | |
| 252 |
no test coverage detected