| 3 | |
| 4 | |
| 5 | def export( |
| 6 | model, quantize: bool = False, opset_version: int = 14, type="onnx", **kwargs |
| 7 | ): |
| 8 | model_scripts = model.export(**kwargs) |
| 9 | export_dir = kwargs.get("output_dir", os.path.dirname(kwargs.get("init_param"))) |
| 10 | os.makedirs(export_dir, exist_ok=True) |
| 11 | |
| 12 | if not isinstance(model_scripts, (list, tuple)): |
| 13 | model_scripts = (model_scripts,) |
| 14 | for m in model_scripts: |
| 15 | m.eval() |
| 16 | if type == "onnx": |
| 17 | _onnx( |
| 18 | m, |
| 19 | quantize=quantize, |
| 20 | opset_version=opset_version, |
| 21 | export_dir=export_dir, |
| 22 | **kwargs, |
| 23 | ) |
| 24 | print("output dir: {}".format(export_dir)) |
| 25 | |
| 26 | return export_dir |
| 27 | |
| 28 | |
| 29 | def _onnx( |