LPRnet中maxpool3d无法成功转ONNX改写

解决ONNX转换问题：YoloV5+LRPnet中maxpool3d的替换策略

最新推荐文章于 2025-05-23 10:16:10 发布

原创

最新推荐文章于 2025-05-23 10:16:10 发布 · 924 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #python #机器学习

文章介绍了在使用YoloV5结合LRPnet进行道路车牌检测和识别的模型训练后，遇到ONNX不支持maxpool3d算子导致模型导出失败的问题。作者提出了通过替换nn.MaxPool3d为自定义的MaxPool3d_muti_batch模块来解决这个问题，该模块利用MaxPool2d和MaxPool1d组合实现类似的功能。经过验证，替换后的代码能够得到相同的输出，从而成功将模型转换为ONNX格式。

使用yolov5+LRPnet进行道路车牌检测和识别，车牌识别模型训练好后使用torch.onnx.export导出的时候遇到了报错，主要问题就是ONNX不支持maxpool3d算子，无法直接进行转换，本文主要针对这一块进行处理

1、LPRnet主干代码

从LPRnet的主代码中可以得知，其中主要有三个地方用到了nn.MaxPool3d，这是本次模型转onnx无法成功的关键，而nn.MaxPool3d只是求取区域内的最大值，不会造成权重的更新，因此只需要更新代码实现方式，获得相对应的输出即可

class LPRNet(nn.Module):
    def __init__(self, lpr_max_len, phase, class_num, dropout_rate):
        super(LPRNet, self).__init__()
        self.phase = phase
        self.lpr_max_len = lpr_max_len
        self.class_num = class_num
        self.backbone = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1),  # 0  [bs,3,24,94] -> [bs,64,22,92]
            nn.BatchNorm2d(num_features=64),  # 1  -> [bs,64,22,92]
            nn.ReLU(),  # 2  -> [bs,64,22,92]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 1, 1)),  # 3  -> [bs,64,20,90]
            small_basic_block(ch_in=64, ch_out=128),  # 4  -> [bs,128,20,90]
            nn.BatchNorm2d(num_features=128),  # 5  -> [bs,128,20,90]
            nn.ReLU(),  # 6  -> [bs,128,20,90]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(2, 1, 2)),  # 7  -> [bs,64,18,44]
            small_basic_block(ch_in=64, ch_out=256),  # 8  -> [bs,256,18,44]
            nn.BatchNorm2d(num_features=256),  # 9  -> [bs,256,18,44]
            nn.ReLU(),  # 10 -> [bs,256,18,44]
            small_basic_block(ch_in=256, ch_out=256),  # 11 -> [bs,256,18,44]
            nn.BatchNorm2d(num_features=256),  # 12 -> [bs,256,18,44]
            nn.ReLU(),  # 13 -> [bs,256,18,44]
            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(4, 1, 2)),  # 14 -> [bs,64,16,21]
            nn.Dropout(dropout_rate),  # 0.5 dropout rate                          # 15 -> [bs,64,16,21]
            nn.Conv2d(in_channels=64, out_channels=256, kernel_size=(1, 4), stride=1),  # 16 -> [bs,256,16,18]
            nn.BatchNorm2d(num_features=256),  # 17 -> [bs,256,16,18]
            nn.ReLU(),  # 18 -> [bs,256,16,18]
            nn.Dropout(dropout_rate),  # 0.5 dropout rate