转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了