pth转ONNX模型结果不一致

转pth模型到ONNX一定要遵循3个原则:

  • 在调用onnx.export 的API前, 一定要时model.eval()
  • 在调用onnx.export 的API前, 一定时刚刚load_state_dict, 之后模型没有任何修改
  • 在调用onnx.export的API前, 模型在正常推理的时候所有处理,转之前都要做

如下面实例所示:

model.load_state_dict(compatible_state_dict, strict=False)
model.eval()
# export to ONNX
torch.onnx.export(
    model,
    x,
    'xxxxx.onnx',
    export_params=True,  # store the trained parameter weights inside the model file 
    opset_version=11,    # the ONNX version to export the model to 
    input_names = ['input'],   # the model's input names 
    output_names = ['xxx1','xxx2'], # the model's output names 
)

坑记录:

为什么我会遇到这个问题呢, 是因为我用的预训练的模型做了一个"moude."字符串删除,如下代码

state_dict = torch.load(path, map_location='cpu')['model']
compatible_state_dict = {}
for k, v in state_dict.items():
    if 'module.' in k:
        compatible_state_dict[k[7:]] = v
    else:
        compatible_state_dict[k] = v

这段代码我在调用onnx.export 的前,没有使用, 就导致导出的onnx推理的结果和原来的pth不一样,也就是犯了第三个准则没有遵循, 所以就GG了

评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

白码思

您的鼓励是我最大的动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值