文字识别领域经典论文回顾第一期：CRNN

class BackBone(nn.Module):
   def __init__(self, inplanes):
       super(BackBone, self).__init__()
       self.inplanes = inplanes
       self.feature_extractor = nn.Sequential(
           nn.Conv2d(self.inplanes, 64, kernel_size=3, stride=1, padding=1),
           nn.MaxPool2d(kernel_size=2, stride=2),
           nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
           nn.ReLU(inplace=True),
           nn.MaxPool2d(kernel_size=2, stride=2),
           nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
           nn.ReLU(inplace=True),
           nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
           nn.ReLU(inplace=True),
         # 第三个maxpooling，注意此处maxpooling的尺寸，可以和论文解读进行对照
           nn.MaxPool2d(kernel_size=(2, 1), stride=(2, 1)),
           nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
           nn.BatchNorm2d(512),
           nn.ReLU(inplace=True),
           nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
           nn.BatchNorm2d(512),
           nn.ReLU(inplace=True),
         # # 第四个maxpooling，同第三个一样
           nn.MaxPool2d(kernel_size=(2, 1), stride=(2, 1)),
           nn.Conv2d(512, 512, kernel_size=(2, 2), stride=1, padding=0),
           nn.ReLU(inplace=True)
      )
​
   def forward(self, inputs):
       return self.feature_extractor(inputs)
     
​
# 这里一般会有一个Reshape layer,对照于论文中的Map-to-Sequence layer,
# 作用是将张量的高度变成1
class ReshapeLayer(nn.Module):
   def __init__(self):
       super(ReshapeLayer, self).__init__()
​
   def forward(self, inputs):
       B, C, H, W = inputs.size()
       inputs = inputs.reshape(B, C, H * W)
       inputs = inputs.permute(0, 2, 1)
       return inputs

class SequenceLayer(nn.Module):
   def __init__(self, num_inputs, num_hiddens):
       super(SequenceLayer, self).__init__()
       self.num_inputs = num_inputs
       self.num_hiddens = num_hiddens
       self.rnn_1 = nn.LSTM(self.num_inputs, self.num_hiddens, bidirectional=True, batch_first=True)
       self.rnn_2 = nn.LSTM(self.num_hiddens, self.num_hiddens, bidirectional=True, batch_first=True)
       self.linear = nn.Linear(self.num_hiddens * 2, self.num_hiddens)
​
   def forward(self, inputs):
       self.rnn_1.flatten_parameters()
       recurrent, _ = self.rnn_1(inputs)
       inputs = self.linear(recurrent)
       self.rnn_2.flatten_parameters()
       recurrent, _ = self.rnn_2(inputs)
       outputs = self.linear(recurrent)
       return outputs

class CTCLoss(nn.Module):
   def __init__(self, params, reduction='mean'):
       super().__init__()
       blank_idx = params.blank_idx
       self.loss_func = torch.nn.CTCLoss(blank=blank_idx, reduction=reduction, zero_infinity=True)
​
   def forward(self, pred, args):
       batch_size = pred.size(0)
       label, label_length = args['targets'], args['targets_lengths']
       pred = pred.log_softmax(2)
       pred = pred.permute(1, 0, 2)
       preds_lengths = torch.tensor([pred.size(0)] * batch_size, dtype=torch.long)
       loss = self.loss_func(pred, label.cuda(), preds_lengths.cuda(), label_length.cuda())
       return loss
     
     
# 最后将上述结构汇总起来
class CRNN(nn.Module):
   def __init__(self, flags):
       super(CRNN, self).__init__()
       self.inplanes = 1 if flags.Global.image_shape[0] == 1 else 3
       self.num_inputs = flags.SeqRNN.input_size
       self.num_hiddens = flags.SeqRNN.hidden_size
       self.converter = CTCLabelConverter(flags)
       self.num_classes = self.converter.char_num
​
       self.feature_extractor = BackBone(self.inplanes)
       self.reshape_layer = ReshapeLayer()
       self.sequence_layer = SequenceLayer(self.num_inputs, self.num_hiddens)
       self.linear_layer = nn.Linear(self.num_hiddens, self.num_classes)
​
   def forward(self, inputs):
       x = self.feature_extractor(inputs)
       x = self.reshape_layer(x)
       x = self.sequence_layer(x)
       outputs = self.linear_layer(x)
​
       return outputs