crnn-pth转换成onnx（pth输入输出变长转换）

注意事项

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)