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了

猜你喜欢

转载自blog.csdn.net/qq_31638535/article/details/128663823