LPRnet pytorch 实现 （参考官方版本）

LPRnet pytorch 实现 （参考官方版本）

之前参考github
https://github.com/sirius-ai/LPRNet_Pytorch
大神LPRnet 代码 用自己生成的数据训练了一个车牌识别模型发现 在训练时候 acc能到95% 与官方论文基本一直 但是测试的时候 效果很差 。在测试的时候 用128 batch size 测试准确率在95% 但是 当单张图片测试的时候 准确率只有80% 一开始以为是 torch .eval() 对BN层处理会有BUG 但是后来找问题在原本代码中
https://github.com/sirius-ai/LPRNet_Pytorch/blob/7c976664b3f3879efabeaff59c7a117e49d5f29e/model/LPRNet.py#L72 这一行 当 batch size 设置为比较大的batch size 时候求的均值 是整个batch 的均值 但是当 为一张图片的时候 均值为图片自己的均值 所以会对预测产生问题。我按照 inter openvion 的官方tensorflow 版本 重新 修改了 模型 模型如下：

// An highlighted block
import torch.nn as nn
import torch

class small_basic_block(nn.Module):
    def __init__(self, ch_in, ch_out):
        super(small_basic_block, self).__init__()
        self.block = nn.Sequential(
            nn.Conv2d(ch_in, ch_out // 4, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(ch_out // 4, ch_out // 4, kernel_size=(3, 1), padding=(1, 0)),
            nn.ReLU(),
            nn.Conv2d(ch_out // 4, ch_out // 4, kernel_size=(1, 3), padding=(0, 1)),
            nn.ReLU(),
            nn.Conv2d(ch_out // 4, ch_out, kernel_size=1),
        )
    def forward(self, x):
        return self.block(x)

class LPRNet(nn.Module):
    def __init__(self,class_num, dropout_rate):
        super(LPRNet, self).__init__()
        self.class_num = class_num
        self.backbone = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1,padding=1), # 0
            nn.BatchNorm2d(num_features=64),
            nn.ReLU(),  # 2
            nn.MaxPool2d(kernel_size=(3, 3), stride=(1, 1)),
            small_basic_block(ch_in=64, ch_out=128),    # *** 4 ***
            nn.BatchNorm2d(num_features=128),
            nn.ReLU(),  # 6
            nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 1)),
            small_basic_block(ch_in=128, ch_out=256),   # 8
            nn.BatchNorm2d(num_features=256),
            nn.ReLU(),  # 10
            small_basic_block(ch_in=256, ch_out=256),   # *** 11 ***
            nn.BatchNorm2d(num_features=256),   # 12
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 1)),  # 14
            nn.Dropout(dropout_rate),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=(4, 1), stride=1),  # 16
            nn.BatchNorm2d(num_features=256),
            nn.ReLU(),  # 18
            nn.Dropout(dropout_rate),
            nn.Conv2d(in_channels=256, out_channels=class_num, kernel_size=(1, 13), stride=1,padding=[0,6]), # 20
            nn.BatchNorm2d(num_features=class_num),
            nn.ReLU(),  # *** 22 ***
        )
        self.connected = nn.Sequential(
            nn.Linear(class_num*88,128),
            nn.ReLU(),
        )
        self.container = nn.Sequential(
            nn.Conv2d(in_channels=128+self.class_num, out_channels=self.class_num, kernel_size=(1, 1), stride=(1, 1)),
            # nn.BatchNorm2d(num_features=self.class_num),
            # nn.ReLU(),
            # nn.Conv2d(in_channels=self.class_num, out_channels=self.lpr_max_len+1, kernel_size=3, stride=2),
            # nn.ReLU(),
        )

    def forward(self, x):


        x = self.backbone(x)
        pattern = x.flatten(1,-1)
        pattern = self.connected(pattern)
        width = x.size()[-1]
        pattern = torch.reshape(pattern,[-1,128,1,1])
        pattern = pattern.repeat(1,1,1,width)
        x = torch.cat([x,pattern],dim=1)
        x = self.container(x)
        logits = x.squeeze(2)


        return logits

def build_lprnet(class_num=66, dropout_rate=0.5):

    Net = LPRNet(class_num, dropout_rate)

    # if phase == "train":
    #     return Net.train()
    # else:
    #     return Net.eval()
    return Net
if __name__ == "__main__":
    from torchsummary import summary
    model = build_lprnet(75,0.5)
    summary(model, (3,24,94), device="cpu")

在我自己相同的数据集 配合上stn 模块 acc在97%