注意事项
1.首先torch1.1.0是没有dynamic_axes这个开关的,所以我升级到了1.14.0。
2.其次如果是变长的话,pth模型找不到input和output的node-name,可以自己随意定义。
3.weights和data必须同样的dtype,可以是cpu也可以是gpu,如果多gpu训练的model,去掉module即可。
4.dynamic_axes找到当前变长的维度,用list或者dict表示,list自动分配名字,很方便。
import torch
import torchvision
import onnx
from crnn.models import crnn as crnn
from collections import OrderedDict
from crnn import keys
from config import ocrModel, LSTMFLAG, GPU
alphabet = keys.alphabetChinese
# LSTMFLAG=True crnn 否则 dense ocr
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256,
1, lstmFlag=True)
state_dict = torch.load(ocrModel,map_location='cpu')
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k.replace('module.', '') # remove `module.`
new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)
model.eval()
input_names = ["input"]
output_names = ["output"]
#batch,channel,h,w
image = torch.zeros(1, 1, 32, 32)
inputs = image
#w->w*32/h
dynamic_axes = {"input":[3],"output":[0]}
torch.onnx.export(model, inputs, './oup_dingchang.onnx',
export_params=True,
verbose=False,
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic_axes,
opset_version=9)