如何CRNN模型训练的pth模型转onnx模型

输入是不定长的文本，通过dynamic_axes={"input": {0:"batch_size",3: "img_w"}}来设置动态输入值

import torch
from lib.utils.utils import strLabelConverter
from lib.models.crnn import CRNN
from collections import OrderedDict
from torch.autograd import Variable
alphabet = '1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ'
nh = 256
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)
model = CRNN(32, 3, len(alphabet) + 1, 256, n_rnn=2, leakyRelu=False)
model.apply(weights_init)
new_state_dict = OrderedDict()
checkpoint = torch.load('models/checkpoint.pth')
for k, v in checkpoint.items():
    print(v)
    name = k
    if k == 'state_dict':
        continue
    if k=='epoch':
        continue
    if k == 'best_acc':
        continue
    new_state_dict[name] = v
model.load_state_dict(checkpoint['state_dict'])
converter = strLabelConverter(alphabet)
# transformer = resizeNormalize((500, 32))
model.eval()
print('Finished loading model!')
print(model)
rnn_input = torch.randn(1, 3, 32, 320)
##################export###############
input_names = ["input"]
output_names = ["out"]
onnx_file_path = 'models/checkpoint.onnx'
torch.onnx.export(model, rnn_input, onnx_file_path, verbose=False,input_names=input_names,output_names=output_names,dynamic_axes={"input": {0:"batch_size",3: "img_w"}}, keep_initializers_as_inputs=True, opset_version=11, export_params=True)
# torch.onnx.export(

如果有些变量名提示没有，那么可以通过该代码段来过滤掉：

for k, v in checkpoint.items():
    print(v)
    name = k
    if k == 'state_dict':
        continue
    if k=='epoch':
        continue
    if k == 'best_acc':
        continue
    new_state_dict[name] = v

如果你训练的单通道图像，那么下面代码中的通道值3要换成1：

model = CRNN(32, 3, len(alphabet) + 1, 256, n_rnn=2, leakyRelu=False)

alphabet值就是你自己工程里训练的字典值。

就这么多吧，还有什么疑问可以留言问，看到会回答的。