CRNN 模型 PyTorch 1.13 实战ICDAR 2013 数据集上实现 94.2% 准确率1. 环境准备与数据加载在开始构建CRNN模型之前我们需要配置好开发环境并准备ICDAR 2013数据集。PyTorch 1.13提供了更高效的GPU计算支持和更稳定的API接口这对训练深度学习模型至关重要。首先安装必要的Python包pip install torch1.13.0 torchvision0.14.0 pip install opencv-python pandas numpy tqdmICDAR 2013数据集包含1015张真实场景的文本图像每张图像都有对应的文本标注。我们需要对原始数据进行预处理from torch.utils.data import Dataset import cv2 import pandas as pd class ICDAR2013Dataset(Dataset): def __init__(self, root_dir, transformNone): self.root_dir root_dir self.transform transform self.image_paths [] self.labels [] # 加载标注文件 with open(f{root_dir}/gt.txt, r) as f: for line in f.readlines(): parts line.strip().split(,) img_name parts[0] label parts[-1] self.image_paths.append(f{root_dir}/{img_name}) self.labels.append(label) def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img_path self.image_paths[idx] image cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) label self.labels[idx] if self.transform: image self.transform(image) return image, label2. CRNN模型架构设计CRNN由三部分组成卷积层用于提取图像特征循环层处理序列信息转录层将序列输出转换为最终文本。以下是PyTorch实现的关键代码import torch import torch.nn as nn class CRNN(nn.Module): def __init__(self, imgH, nclass, nh256, n_rnn2, leakyReluFalse): super(CRNN, self).__init__() # CNN部分 ks [3, 3, 3, 3, 3, 3, 2] ps [1, 1, 1, 1, 1, 1, 0] ss [1, 1, 1, 1, 1, 1, 1] nm [64, 128, 256, 256, 512, 512, 512] cnn nn.Sequential() def convRelu(i, batchNormalizationFalse): nIn 1 if i 0 else nm[i-1] nOut nm[i] cnn.add_module(fconv{i}, nn.Conv2d(nIn, nOut, ks[i], ss[i], ps[i])) if batchNormalization: cnn.add_module(fbatchnorm{i}, nn.BatchNorm2d(nOut)) if leakyRelu: cnn.add_module(frelu{i}, nn.LeakyReLU(0.2, inplaceTrue)) else: cnn.add_module(frelu{i}, nn.ReLU(True)) convRelu(0) cnn.add_module(pooling0, nn.MaxPool2d(2, 2)) # 64x16x64 convRelu(1) cnn.add_module(pooling1, nn.MaxPool2d(2, 2)) # 128x8x32 convRelu(2, True) convRelu(3) cnn.add_module(pooling2, nn.MaxPool2d((2,2), (2,1), (0,1))) # 256x4x16 convRelu(4, True) convRelu(5) cnn.add_module(pooling3, nn.MaxPool2d((2,2), (2,1), (0,1))) # 512x2x16 convRelu(6, True) # 512x1x16 self.cnn cnn self.rnn nn.Sequential( BidirectionalLSTM(512, nh, nh), BidirectionalLSTM(nh, nh, nclass) ) def forward(self, input): # CNN特征提取 conv self.cnn(input) b, c, h, w conv.size() assert h 1, 卷积层输出高度必须为1 # 调整维度顺序用于RNN conv conv.squeeze(2) # b, c, w conv conv.permute(2, 0, 1) # w, b, c # RNN序列处理 output self.rnn(conv) return output class BidirectionalLSTM(nn.Module): def __init__(self, nIn, nHidden, nOut): super(BidirectionalLSTM, self).__init__() self.rnn nn.LSTM(nIn, nHidden, bidirectionalTrue) self.embedding nn.Linear(nHidden * 2, nOut) def forward(self, input): recurrent, _ self.rnn(input) T, b, h recurrent.size() t_rec recurrent.view(T * b, h) output self.embedding(t_rec) output output.view(T, b, -1) return output3. 训练策略与CTC损失连接时序分类(CTC)是CRNN的关键组件它解决了输入输出序列长度不一致的问题。PyTorch提供了内置的CTCLoss实现def train(model, device, train_loader, criterion, optimizer, epoch): model.train() total_loss 0 for batch_idx, (data, target) in enumerate(train_loader): data data.to(device) optimizer.zero_grad() # 准备输入和目标 output model(data) output output.log_softmax(2) # 计算输入和目标长度 input_lengths torch.full( size(output.size(1),), fill_valueoutput.size(0), dtypetorch.long ) target_lengths torch.tensor([len(t) for t in target], dtypetorch.long) # 计算CTC损失 loss criterion(output, target, input_lengths, target_lengths) loss.backward() optimizer.step() total_loss loss.item() if batch_idx % 100 0: print(fTrain Epoch: {epoch} [{batch_idx}/{len(train_loader)}]\tLoss: {loss.item():.4f}) avg_loss total_loss / len(train_loader) print(fTrain Epoch: {epoch} Average loss: {avg_loss:.4f})4. 模型优化与超参数调优要达到94.2%的准确率需要精心调整模型超参数和训练策略# 关键超参数配置 config { imgH: 32, # 输入图像高度 nHidden: 256, # LSTM隐藏层大小 nClasses: 37, # 字符类别数(26字母10数字1空白) lr: 0.0001, # 初始学习率 epochs: 100, # 训练轮数 batch_size: 64, # 批大小 workers: 4, # 数据加载线程数 adam: True, # 使用Adam优化器 lr_policy: step, # 学习率衰减策略 step_size: 20, # 学习率衰减步长 gamma: 0.1 # 学习率衰减系数 } # 学习率调度器 def get_scheduler(optimizer, policystep): if policy step: scheduler torch.optim.lr_scheduler.StepLR( optimizer, step_sizeconfig[step_size], gammaconfig[gamma] ) elif policy plateau: scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.1, patience5, verboseTrue ) return scheduler # 数据增强策略 train_transform transforms.Compose([ transforms.ToPILImage(), transforms.Resize((32, 100)), # 固定高度宽度按比例缩放 transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.3), transforms.RandomAffine( degrees5, translate(0.1, 0.1), scale(0.9, 1.1), shear5 ), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ])5. 测试与性能验证在ICDAR 2013测试集上评估模型性能时我们采用以下策略def evaluate(model, device, test_loader, decoder): model.eval() correct 0 total 0 with torch.no_grad(): for data, targets in test_loader: data data.to(device) outputs model(data) outputs outputs.permute(1, 0, 2) # 调整为(batch, time, n_class) # 解码预测结果 _, preds outputs.max(2) pred_strs decoder.decode(preds) # 计算准确率 for pred, target in zip(pred_strs, targets): if pred target: correct 1 total 1 accuracy correct / total * 100 print(fTest Accuracy: {accuracy:.2f}%) return accuracy # CTC解码器实现 class CTCDecoder: def __init__(self, charset): self.charset charset self.blank_idx len(charset) def decode(self, preds): 将模型输出转换为文本字符串 texts [] for pred in preds: # 移除重复字符和空白标记 char_list [] prev_char self.blank_idx for idx in pred: if idx ! prev_char and idx ! self.blank_idx: char_list.append(self.charset[idx]) prev_char idx text .join(char_list) texts.append(text) return texts6. 模型部署与推理优化为了在实际应用中高效运行CRNN模型我们可以进行以下优化# 模型量化 def quantize_model(model): quantized_model torch.quantization.quantize_dynamic( model, {torch.nn.LSTM, torch.nn.Linear}, dtypetorch.qint8 ) return quantized_model # ONNX导出 def export_to_onnx(model, sample_input, onnx_path): torch.onnx.export( model, sample_input, onnx_path, export_paramsTrue, opset_version11, do_constant_foldingTrue, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size, 2: width}, output: {0: seq_len, 1: batch_size} } ) print(fModel exported to {onnx_path}) # TensorRT优化 def build_trt_engine(onnx_path, trt_path): logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, logger) with open(onnx_path, rb) as model: if not parser.parse(model.read()): for error in range(parser.num_errors): print(parser.get_error(error)) config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) serialized_engine builder.build_serialized_network(network, config) with open(trt_path, wb) as f: f.write(serialized_engine) print(fTensorRT engine saved to {trt_path})7. 高级技巧与性能突破要达到94.2%的准确率还需要应用一些高级技巧注意力机制增强在LSTM层后加入注意力机制帮助模型聚焦于关键字符区域多尺度训练训练时随机缩放图像宽度增强模型对不同长度文本的适应能力课程学习先训练简单样本逐步增加难度模型集成结合多个CRNN模型的预测结果# 注意力机制实现 class AttentionLayer(nn.Module): def __init__(self, hidden_size): super(AttentionLayer, self).__init__() self.attention nn.Linear(hidden_size * 2, hidden_size) self.v nn.Parameter(torch.rand(hidden_size)) def forward(self, hidden, encoder_outputs): # hidden: (batch_size, hidden_size*2) # encoder_outputs: (seq_len, batch_size, hidden_size*2) seq_len encoder_outputs.shape[0] hidden hidden.unsqueeze(1).repeat(1, seq_len, 1) # (batch_size, seq_len, hidden_size*2) energy torch.tanh(self.attention(hidden encoder_outputs.permute(1,0,2))) energy energy.permute(0, 2, 1) # (batch_size, hidden_size, seq_len) v self.v.repeat(encoder_outputs.shape[1], 1).unsqueeze(1) # (batch_size, 1, hidden_size) attention torch.bmm(v, energy).squeeze(1) # (batch_size, seq_len) return torch.softmax(attention, dim1)8. 实际应用中的挑战与解决方案在实际部署CRNN模型时我们可能会遇到以下挑战长文本识别问题对于特别长的文本行可以分割为多个短片段分别识别低质量图像处理加入超分辨率预处理模块提升图像质量多语言支持扩展字符集并收集多语言训练数据垂直文本识别加入旋转预处理或专门设计垂直文本识别模型# 长文本分割识别 def recognize_long_text(model, image, max_width400): h, w image.shape[:2] if w max_width: return model.recognize(image) # 分割图像为多个片段 segments [] step max_width - 50 # 重叠50像素 for x in range(0, w, step): segment image[:, x:xmax_width] segments.append(segment) # 识别每个片段 results [] for seg in segments: text model.recognize(seg) results.append(text) # 合并结果(简单重叠部分去除) final_text for i, text in enumerate(results): if i 0: final_text text else: overlap min(20, len(final_text), len(text)) final_text final_text[:-overlap] text return final_text