Spaces:
Running
Running
| import os | |
| import sys | |
| __dir__ = os.path.dirname(os.path.abspath(__file__)) | |
| sys.path.append(__dir__) | |
| sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '..'))) | |
| import torch | |
| from openrec.modeling import build_model | |
| from openrec.postprocess import build_post_process | |
| from tools.engine import Config | |
| from tools.infer_rec import build_rec_process | |
| from tools.utility import ArgsParser | |
| from tools.utils.ckpt import load_ckpt | |
| from tools.utils.logging import get_logger | |
| def to_onnx(model, dummy_input, dynamic_axes, sava_path='model.onnx'): | |
| input_axis_name = ['batch_size', 'channel', 'in_width', 'int_height'] | |
| output_axis_name = ['batch_size', 'channel', 'out_width', 'out_height'] | |
| torch.onnx.export( | |
| model.to('cpu'), | |
| dummy_input, | |
| sava_path, | |
| input_names=['input'], | |
| output_names=['output'], # the model's output names | |
| dynamic_axes={ | |
| 'input': {axis: input_axis_name[axis] | |
| for axis in dynamic_axes}, | |
| 'output': {axis: output_axis_name[axis] | |
| for axis in dynamic_axes}, | |
| }, | |
| ) | |
| def export_single_model(model: torch.nn.Module, _cfg, export_dir, | |
| export_config, logger, type): | |
| for layer in model.modules(): | |
| if hasattr(layer, 'rep') and not getattr(layer, 'is_repped'): | |
| layer.rep() | |
| os.makedirs(export_dir, exist_ok=True) | |
| export_cfg = {'PostProcess': _cfg['PostProcess']} | |
| export_cfg['Transforms'] = build_rec_process(_cfg) | |
| cfg.save(os.path.join(export_dir, 'config.yaml'), export_cfg) | |
| dummy_input = torch.randn(*export_config['export_shape'], device='cpu') | |
| if type == 'script': | |
| save_path = os.path.join(export_dir, 'model.pt') | |
| trace_model = torch.jit.trace(model, dummy_input, strict=False) | |
| torch.jit.save(trace_model, save_path) | |
| elif type == 'onnx': | |
| save_path = os.path.join(export_dir, 'model.onnx') | |
| to_onnx(model, dummy_input, export_config.get('dynamic_axes', []), | |
| save_path) | |
| else: | |
| raise NotImplementedError | |
| logger.info(f'finish export model to {save_path}') | |
| def main(cfg, type): | |
| _cfg = cfg.cfg | |
| logger = get_logger() | |
| global_config = _cfg['Global'] | |
| export_config = _cfg['Export'] | |
| # build post process | |
| post_process_class = build_post_process(_cfg['PostProcess']) | |
| char_num = len(getattr(post_process_class, 'character')) | |
| cfg['Architecture']['Decoder']['out_channels'] = char_num | |
| model = build_model(_cfg['Architecture']) | |
| load_ckpt(model, _cfg) | |
| model.eval() | |
| export_dir = export_config.get('export_dir', '') | |
| if not export_dir: | |
| export_dir = os.path.join(global_config.get('output_dir', 'output'), | |
| 'export') | |
| if _cfg['Architecture']['algorithm'] in ['Distillation' | |
| ]: # distillation model | |
| _cfg['PostProcess'][ | |
| 'name'] = post_process_class.__class__.__base__.__name__ | |
| for model_name in model.model_list: | |
| sub_model_save_path = os.path.join(export_dir, model_name) | |
| export_single_model( | |
| model.model_list[model_name], | |
| _cfg, | |
| sub_model_save_path, | |
| export_config, | |
| logger, | |
| type, | |
| ) | |
| else: | |
| export_single_model(model, _cfg, export_dir, export_config, logger, | |
| type) | |
| def parse_args(): | |
| parser = ArgsParser() | |
| parser.add_argument('--type', | |
| type=str, | |
| default='onnx', | |
| help='type of export') | |
| args = parser.parse_args() | |
| return args | |
| if __name__ == '__main__': | |
| FLAGS = parse_args() | |
| cfg = Config(FLAGS.config) | |
| FLAGS = vars(FLAGS) | |
| opt = FLAGS.pop('opt') | |
| cfg.merge_dict(FLAGS) | |
| cfg.merge_dict(opt) | |
| main(cfg, FLAGS['type']) | |