| 35 | |
| 36 | |
| 37 | def export_single_model(model: torch.nn.Module, _cfg, export_dir, |
| 38 | export_config, logger, type): |
| 39 | for layer in model.modules(): |
| 40 | if hasattr(layer, 'rep') and not getattr(layer, 'is_repped'): |
| 41 | layer.rep() |
| 42 | os.makedirs(export_dir, exist_ok=True) |
| 43 | |
| 44 | export_cfg = {'PostProcess': _cfg['PostProcess']} |
| 45 | export_cfg['Transforms'] = build_rec_process(_cfg) |
| 46 | |
| 47 | cfg.save(os.path.join(export_dir, 'config.yaml'), export_cfg) |
| 48 | |
| 49 | dummy_input = torch.randn(*export_config['export_shape'], device='cpu') |
| 50 | if type == 'script': |
| 51 | save_path = os.path.join(export_dir, 'model.pt') |
| 52 | trace_model = torch.jit.trace(model, dummy_input, strict=False) |
| 53 | torch.jit.save(trace_model, save_path) |
| 54 | elif type == 'onnx': |
| 55 | save_path = os.path.join(export_dir, 'model.onnx') |
| 56 | to_onnx(model, dummy_input, export_config.get('dynamic_axes', []), |
| 57 | save_path) |
| 58 | else: |
| 59 | raise NotImplementedError |
| 60 | logger.info(f'finish export model to {save_path}') |
| 61 | |
| 62 | |
| 63 | def main(cfg, type): |