pytorch训练的pt模型转换为onnx(nn.DataParallel()、model、model.state_dict())

pt转onnx流程与常见问题

pt转onnx流程

1、读取pt模型文件,文件既可以是torch.save(model,path)整体保存的模型,也可以是保存的字典文件。

// An highlighted block
def load_model(model, model_path):
    checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)

    try:
        state_dict_ = checkpoint["state_dict"]
        state_dict = {}
    except:
        state_dict_ = checkpoint
        state_dict = {}
    # convert data_parallal to model
    for k in state_dict_:

        if k.startswith('module') and not k.startswith('module_list'):
            state_dict[k[7:]] = state_dict_[k]
        else:
            state_dict[k] = state_dict_[k]
    model_state_dict = model.state_dict()

    # check loaded parameters and created model parameters
    msg = 'If you see this, your model does not fully load the ' + \
          'pre-trained weight. Please make sure ' + \
          'you have correctly specified --arch xxx ' + \
          'or set the correct --num_classes for your own dataset.'
    for k in state_dict:
        if k in model_state_dict:
            if state_dict[k].shape != model_state_dict[k].shape:
                print('Skip loading parameter {}, required shape{}, ' \
                      'loaded shape{}. {}'.format(
                    k, model_state_dict[k].shape, state_dict[k].shape, msg))
                state_dict[k] = model_state_dict[k]
        else:
            print('Drop parameter {}.'.format(k) + msg)
    for k in model_state_dict:
        if not (k in state_dict):
            print('No param {}.'.format(k) + msg)
            state_dict[k] = model_state_dict[k]
    model.load_state_dict(state_dict, strict=False)
    return model

2、torch.onnx.export()函数转换

checkpoint = torch.load(weights, map_location=torch.device('cpu'))#weight为模型路径
model.load_state_dict(checkpoint["model_state"])
im = torch.zeros([1, 3, 1080, 1920])#输入的尺寸,根据自己的情况写死
onnx_path = ''#onnx模型路径
torch.onnx.export(model, im, onnx_path,verbose=False,opset_version = 12,do_constant_folding = False,input_names=['images'],output_names=['output'])
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值