深度学习实战:基于MobileNetV2的HWDB1.1手写汉字识别系统构建
1. 从零开始:为什么选择MobileNetV2和HWDB1.1?
大家好,我是老张,在AI和嵌入式设备上折腾了十多年。今天想和大家聊聊一个特别有意思,也特别有挑战性的实战项目:用MobileNetV2这个轻量级网络,来搞定HWDB1.1手写汉字识别。你可能会问,汉字识别不是老生常谈了吗?干嘛还要专门写一篇?嘿,这你可就问到点子上了。
我见过太多教程,一上来就甩给你一个ResNet或者VGG的代码,跑在MNIST这种简单数据集上,准确率轻松99%+,看着很唬人。但一旦你拿到像HWDB1.1这样包含3755个类别、超过一百万张图片的真实、复杂数据集,那些“玩具”项目里的经验就完全不够用了。模型瞬间变得又大又慢,训练几天几夜,部署到资源有限的设备上更是痴人说梦。这就是我们这个项目的核心价值:在保证高精度的前提下,追求极致的效率。
所以,我选择了MobileNetV2。它不是性能最强的,但绝对是“性价比”最高的选手之一。它用了倒残差结构和线性瓶颈层,能在参数和计算量大幅减少的同时,保持不错的特征提取能力。简单说,它就是为移动端和嵌入式设备而生的。而HWDB1.1数据集,由中科院自动化所出品,是国内手写汉字识别领域的权威基准,数据量大、书写风格多样,非常考验模型的泛化能力。把这两者结合起来,构建一个从数据预处理、模型训练到优化部署的完整流水线,这才是一个有工业级价值的实战项目。
这个项目适合谁呢?如果你是深度学习初学者,想找一个有足够深度的项目来练手,超越简单的猫狗分类;如果你是移动端AI开发者,正在为模型裁剪和加速发愁;或者你单纯对OCR技术感兴趣,想了解如何处理超多类别的分类问题——那么,跟着我一步步走下来,你收获的将不仅仅是一段能跑的代码,更是一套解决复杂实际问题的工程化思维。好了,闲话少说,我们撸起袖子开干。
2. 搞定“硬骨头”:HWDB1.1数据集的解析与预处理
拿到HWDB1.1数据集,很多朋友第一步就卡住了。官网下载下来的不是常见的jpg或png,而是一种叫做.gnt的二进制文件。这算是我们遇到的第一个“坑”,但别怕,拆解它其实很有意思。
2.1 理解GNT文件结构与解析原理
.gnt文件是当年为了方便存储和传输大量手写样本而设计的格式。你可以把它想象成一个容器,里面按顺序打包了成千上万张图片和对应的标签。它的结构很有规律:每张图片都由一个文件头和图片像素数据连续组成。
文件头固定10个字节,包含了这张图片的关键信息。我来给你拆解一下:前4个字节(按小端序)拼起来,表示从文件头开始到这张图片结束的总字节数;接着2个字节是汉字的国标码;最后4个字节分别代表图片的宽和高。理解这个结构是解析的关键。后面紧跟着的,就是宽度*高度个字节的像素数据,直接按顺序读出来,reshape一下,一张灰度图就出来了。
原始文章里给出了解析代码,但我想结合我的踩坑经验,给你一个更健壮、更清晰的版本。这里最大的一个坑是:HWDB1.1的训练集压缩包解压后是.alz格式,需要二次解压。在Linux下可以用unalz命令,Windows则需要找专门的解压工具(比如Bandizip的新版本就支持)。
2.2 实战代码:稳健的GNT转PNG流水线
理解了原理,我们来看代码。解析的核心是一个生成器函数,它负责从一个.gnt文件里源源不断地读出图片和标签。这里要特别注意字节序的处理和文件读取的边界判断,否则很容易读乱套。
import os
import numpy as np
import struct
from PIL import Image
import pickle
def read_gnt_file(gnt_path):
"""
从单个.gnt文件中读取所有(图像,标签)对
"""
with open(gnt_path, 'rb') as f:
while True:
# 1. 读取10字节的头部
header = np.fromfile(f, dtype=np.uint8, count=10)
if header.size == 0: # 文件结束
break
# 2. 解析头部信息(注意字节序,这里是小端)
sample_size = header[0] + (header[1]<<8) + (header[2]<<16) + (header[3]<<24)
tagcode = header[5] + (header[4]<<8) # 国标码
width = header[6] + (header[7]<<8)
height = header[8] + (header[9]<<8)
# 3. 安全性校验:理论上 10 + width*height 应等于 sample_size
# 有些损坏的文件这里会出错,加个判断更稳妥
if 10 + width * height != sample_size:
# 计算实际应读取的像素数,防止错位
actual_pixel_count = sample_size - 10
# 简单处理:调整宽度或高度,这里选择按比例调整宽度
width = actual_pixel_count // height
# 你也可以选择跳过这张图,我建议打印个警告
print(f"警告: 文件 {gnt_path} 中样本头信息不一致,已调整宽度为 {width}")
# 4. 读取图像数据
image_data = np.fromfile(f, dtype=np.uint8, count=width*height).reshape((height, width))
# 5. 将国标码转换为Unicode汉字
# GB2312编码是双字节,struct.pack把它打包后再解码
try:
character = struct.pack('>H', tagcode).decode('gb2312', errors='ignore')
# 忽略解码错误,极少数情况可能遇到非法编码
except:
continue # 跳过无法解码的样本
yield image_data, character
def convert_gnt_to_images(gnt_dir, output_root_dir, mode='train'):
"""
将某个目录下所有.gnt文件转换为按标签分类的png图片
mode: 'train' 或 'test',用于创建子目录
"""
# 首先,遍历一次所有文件,收集所有出现的汉字,构建标签字典
print(f"正在扫描 {gnt_dir} 以构建字符集...")
char_set = set()
for filename in os.listdir(gnt_dir):
if filename.endswith('.gnt'):
file_path = os.path.join(gnt_dir, filename)
for _, char in read_gnt_file(file_path):
if char: # 确保字符有效
char_set.add(char)
char_list = sorted(list(char_set)) # 排序以保证每次生成的顺序一致
char_to_idx = {char: idx for idx, char in enumerate(char_list)}
idx_to_char = {idx: char for char, idx in char_to_idx.items()}
# 保存标签映射字典,后续训练和预测都要用到
dict_save_path = os.path.join(output_root_dir, f'char_dict_{mode}.pkl')
with open(dict_save_path, 'wb') as f:
pickle.dump({'char_to_idx': char_to_idx, 'idx_to_char': idx_to_char}, f)
print(f"字符字典已保存至 {dict_save_path}, 共 {len(char_list)} 个唯一字符。")
# 创建输出目录(按标签ID建立子文件夹)
image_output_dir = os.path.join(output_root_dir, mode)
os.makedirs(image_output_dir, exist_ok=True)
for idx in range(len(char_list)):
os.makedirs(os.path.join(image_output_dir, f'{idx:05d}'), exist_ok=True)
# 第二次遍历,正式转换并保存图片
print(f"开始转换图片并保存至 {image_output_dir} ...")
counter = 0
for filename in os.listdir(gnt_dir):
if filename.endswith('.gnt'):
file_path = os.path.join(gnt_dir, filename)
for img_array, char in read_gnt_file(file_path):
if char not in char_to_idx:
continue # 跳过不在字典中的字符(理论上不会发生)
label_idx = char_to_idx[char]
# 使用PIL将numpy数组保存为PNG
img = Image.fromarray(img_array).convert('L') # 转为灰度图
save_path = os.path.join(image_output_dir, f'{label_idx:05d}', f'{counter:08d}.png')
img.save(save_path)
counter += 1
if counter % 5000 == 0:
print(f"已处理 {counter} 张图片...")
print(f"转换完成!共处理 {counter} 张图片。")
return len(char_list)
# 使用示例
if __name__ == '__main__':
# 请根据你的实际路径修改
train_gnt_dir = '/path/to/HWDB1.1trn_gnt' # 解压后的训练集gnt文件夹
test_gnt_dir = '/path/to/HWDB1.1tst_gnt' # 解压后的测试集gnt文件夹
output_root = './HWDB1.1_dataset'
num_train_classes = convert_gnt_to_images(train_gnt_dir, output_root, mode='train')
num_test_classes = convert_gnt_to_images(test_gnt_dir, output_root, mode='test')
print(f"训练集类别数: {num_train_classes}, 测试集类别数: {num_test_classes}")
这段代码我做了几个关键优化:一是增加了文件头校验,防止因个别损坏文件导致整个解析中断;二是将构建字符字典和转换图片分成了两步,逻辑更清晰,并且保存了映射关系文件;三是图片命名用了8位数字填充,避免排序混乱。预处理这一步虽然繁琐,但地基打好了,后面训练才能顺风顺水。
3. 模型核心:深入理解并实现MobileNetV2
数据准备好了,接下来就是重头戏:模型。为什么不用更火的Vision Transformer或者更大的CNN?原因很简单:我们要考虑落地。在手机或者边缘计算盒子上跑一个3755类的分类模型,速度和模型大小是硬指标。MobileNetV2就是在这个平衡木上跳舞的高手。
3.1 拆解MobileNetV2的“独门秘籍”
MobileNetV2的核心创新是倒残差结构和线性瓶颈层。这俩词听起来挺唬人,我用个类比你就明白了。传统的残差块(比如ResNet)是“胖-瘦-胖”的结构:先用一个1x1卷积把通道数压缩(变瘦),再用3x3卷积提取特征,最后再用1x1卷积把通道数扩张回去(变胖)。MobileNetV2反其道而行之,是“瘦-胖-瘦”:先升维,在更高维的空间里做深度可分离卷积(这是MobileNet系列的精髓,能大幅减少计算量),最后再降维。
你可能会问,先升维不是增加了计算量吗?妙就妙在,它升维之后使用的是深度可分离卷积。这种卷积把标准卷积拆成两步:先对每个输入通道单独做卷积(深度卷积),再用1x1卷积(逐点卷积)组合通道。计算量能降到原来的差不多十分之一。所以,总体算下来,这个“倒残差”结构依然比标准卷积块轻量得多。
那线性瓶颈层又是什么?在最后一个1x1卷积后,MobileNetV2去掉了ReLU激活函数,改用线性输出。这是因为作者发现,在低维空间使用ReLU会破坏特征信息。想象一下,把高维特征压缩到一个小空间里,再用ReLU这种会“掐掉”负值的函数,信息损失就太严重了。去掉ReLU,保留完整的线性变换,特征信息就保住了。
3.2 手把手实现MobileNetV2网络
光说不练假把式,我们对照着论文里的结构图,用PyTorch把它实现出来。我在这里会实现一个可配置的版本,方便你调整宽度乘数(width multiplier)来进一步压缩模型。
import torch
import torch.nn as nn
import torch.nn.functional as F
class ConvBNReLU(nn.Sequential):
"""一个标准的卷积+BN+ReLU组合块,经常用到,单独提出来"""
def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
padding = (kernel_size - 1) // 2
super(ConvBNReLU, self).__init__(
nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
nn.BatchNorm2d(out_planes),
nn.ReLU6(inplace=True) # MobileNetV2中使用ReLU6作为激活函数
)
class InvertedResidual(nn.Module):
"""倒残差块,MobileNetV2的核心模块"""
def __init__(self, inp, oup, stride, expand_ratio):
"""
Args:
inp: 输入通道数
oup: 输出通道数
stride: 步长
expand_ratio: 扩张倍数(即中间层的通道数是输入的多少倍)
"""
super(InvertedResidual, self).__init__()
self.stride = stride
assert stride in [1, 2]
hidden_dim = int(round(inp * expand_ratio))
# 是否使用快捷连接:只有当步长为1且输入输出通道数相同时才用
self.use_res_connect = self.stride == 1 and inp == oup
layers = []
if expand_ratio != 1:
# 第一阶段:逐点卷积进行升维 (1x1)
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
layers.extend([
# 第二阶段:深度可分离卷积 (3x3)
ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
# 第三阶段:逐点卷积进行降维,注意这里没有ReLU!
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
nn.BatchNorm2d(oup),
])
self.conv = nn.Sequential(*layers)
def forward(self, x):
if self.use_res_connect:
return x + self.conv(x)
else:
return self.conv(x)
class MobileNetV2(nn.Module):
def __init__(self, num_classes=3755, width_mult=1.0, input_size=32):
"""
Args:
num_classes: 分类数,我们这里是3755
width_mult: 宽度乘数,用于整体缩放通道数,默认1.0是原版
input_size: 输入图像尺寸,HWDB1.1我们resize到32x32
"""
super(MobileNetV2, self).__init__()
# 根据输入尺寸调整最后的平均池化核大小
last_channel = 1280
if input_size == 32:
final_pool_size = 4
elif input_size == 64:
final_pool_size = 8
else:
# 其他尺寸可以按比例估算,简单起见这里固定
final_pool_size = 4
print(f"警告: 输入尺寸{input_size}非标准值,池化核暂设为{final_pool_size}")
# 初始卷积层
input_channel = int(32 * width_mult)
self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel
features = [ConvBNReLU(3, input_channel, stride=2 if input_size>=64 else 1)] # 小图不用下采样太快
# 倒残差块的配置表: [扩张倍数t, 输出通道c, 重复次数n, 步长s]
inverted_residual_setting = [
# t, c, n, s
[1, 16, 1, 1],
[6, 24, 2, 2],
[6, 32, 3, 2],
[6, 64, 4, 2],
[6, 96, 3, 1],
[6, 160, 3, 2],
[6, 320, 1, 1],
]
# 根据配置表构建网络主体
for t, c, n, s in inverted_residual_setting:
output_channel = int(c * width_mult)
for i in range(n):
stride = s if i == 0 else 1 # 每个stage只有第一个block进行下采样
features.append(InvertedResidual(input_channel, output_channel, stride, expand_ratio=t))
input_channel = output_channel
# 最后的1x1卷积层
features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1))
self.features = nn.Sequential(*features)
# 分类器
self.classifier = nn.Sequential(
nn.Dropout(0.2), # 原论文在分类前使用了Dropout
nn.Linear(self.last_channel, num_classes),
)
# 权重初始化
self._initialize_weights()
def forward(self, x):
x = self.features(x)
# 全局平均池化
x = F.adaptive_avg_pool2d(x, (1, 1))
x = torch.flatten(x, 1)
x = self.classifier(x)
return x
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out')
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.BatchNorm2d):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.zeros_(m.bias)
# 快速测试一下模型
if __name__ == '__main__':
# 创建一个针对32x32输入,3755类的MobileNetV2
model = MobileNetV2(num_classes=3755, input_size=32)
print(model)
# 计算参数量
total_params = sum(p.numel() for p in model.parameters())
print(f"模型总参数量: {total_params / 1e6:.2f} M")
# 模拟一个输入
dummy_input = torch.randn(2, 3, 32, 32)
output = model(dummy_input)
print(f"输入形状: {dummy_input.shape}, 输出形状: {output.shape}")
这个实现比原始文章里的更模块化,也更贴近原论文。我增加了width_mult参数,你可以通过调小它(比如0.5)来得到一个更轻量的模型,虽然精度可能会略有下降,但在资源受限的场景下非常有用。另外,根据输入尺寸动态调整初始步长和最终池化核大小,让模型适应性更强。
4. 训练策略与技巧:让轻量模型也能学得好
模型搭好了,直接开训?别急,对于HWDB1.1这种类别极度不平衡(有的字样本多,有的字样本少)、且模型容量相对较小的任务,训练策略至关重要。用对方法,MobileNetV2也能在轻量级选手中脱颖而出。
4.1 数据加载与增强:给模型“喂”好数据
PyTorch的DataLoader和torchvision.transforms是我们的好帮手。对于手写汉字,我们需要设计合适的增强策略。不能乱增强,比如随机垂直翻转汉字,那可能就变成另一个字了。水平翻转也要谨慎。我推荐的增强组合是:小角度的随机旋转(比如±10度)、轻微的随机仿射变换(模拟书写抖动)、以及标准化。
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
import os
from PIL import Image
class HWDBDataset(Dataset):
"""自定义HWDB1.1数据集类"""
def __init__(self, root_dir, transform=None, mode='train'):
self.root_dir = root_dir
self.transform = transform
self.mode = mode
self.data = []
self.labels = []
# 读取之前保存的标签字典
dict_path = os.path.join(root_dir, f'char_dict_{mode}.pkl')
with open(dict_path, 'rb') as f:
char_dict = pickle.load(f)
self.idx_to_char = char_dict['idx_to_char']
self.char_to_idx = char_dict['char_to_idx']
# 遍历所有类别文件夹,收集图片路径和标签
class_dirs = sorted([d for d in os.listdir(os.path.join(root_dir, mode)) if os.path.isdir(os.path.join(root_dir, mode, d))])
for label_idx_str in class_dirs:
label_idx = int(label_idx_str)
class_path = os.path.join(root_dir, mode, label_idx_str)
for img_name in os.listdir(class_path):
if img_name.endswith('.png'):
self.data.append(os.path.join(class_path, img_name))
self.labels.append(label_idx)
print(f"{mode}集加载完成,共 {len(self.data)} 张图片,{len(class_dirs)} 个类别。")
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
img_path = self.data[idx]
label = self.labels[idx]
# 用PIL打开图片,确保是灰度图(单通道)
image = Image.open(img_path).convert('L')
# 灰度图转RGB三通道,以适应预训练模型或某些网络结构要求
image = image.convert('RGB')
if self.transform:
image = self.transform(image)
return image, label
# 定义训练和测试的数据增强
train_transform = transforms.Compose([
transforms.RandomAffine(degrees=10, translate=(0.05, 0.05), scale=(0.9, 1.1), shear=5),
transforms.RandomHorizontalFlip(p=0.1), # 小概率水平翻转,汉字基本对称
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) # 归一化到[-1, 1]
])
test_transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
# 创建数据集和数据加载器
train_dataset = HWDBDataset('./HWDB1.1_dataset', transform=train_transform, mode='train')
test_dataset = HWDBDataset('./HWDB1.1_dataset', transform=test_transform, mode='test')
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=4, pin_memory=True)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=4, pin_memory=True)
这里我用了RandomAffine来做仿射变换,模拟书写时的轻微形变,这是提升模型泛化能力的关键。num_workers和pin_memory可以加速数据加载,尤其是在GPU训练时。
4.2 损失函数与优化器调参:针对多分类的“组合拳”
对于3755类的分类,交叉熵损失nn.CrossEntropyLoss是标准选择。但优化器的选择和学习率调度是门艺术。我试过很多组合,对于MobileNetV2,SGD with Momentum配合Cosine Annealing LR Scheduler效果非常稳定。
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
def create_model_and_optimizer(num_classes, device, learning_rate=0.05, weight_decay=4e-5):
"""创建模型、优化器、损失函数和学习率调度器"""
model = MobileNetV2(num_classes=num_classes, input_size=32).to(device)
# 损失函数
criterion = nn.CrossEntropyLoss(label_smoothing=0.1) # 标签平滑,防止过拟合
# 优化器:SGD with Momentum 和权重衰减
optimizer = optim.SGD(model.parameters(),
lr=learning_rate,
momentum=0.9,
weight_decay=weight_decay,
nesterov=True) # 使用Nesterov动量
# 学习率调度:热身+余弦退火
# 先线性预热5个epoch,再余弦下降
warmup_epochs = 5
total_epochs = 100
def warmup_cosine_scheduler(epoch):
if epoch < warmup_epochs:
# 线性预热
return (epoch + 1) / warmup_epochs
else:
# 余弦退火
progress = (epoch - warmup_epochs) / (total_epochs - warmup_epochs)
return 0.5 * (1.0 + math.cos(math.pi * progress))
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=warmup_cosine_scheduler)
return model, criterion, optimizer, scheduler
这里我用了几个小技巧:标签平滑让模型不那么“自信”,减轻过拟合;Nesterov动量让SGD收敛更快更稳;学习率预热在训练初期用小学习率,避免梯度震荡,然后再用余弦退火平滑地降低学习率。这些技巧对训练深度网络,尤其是轻量网络,帮助巨大。
4.3 训练循环与可视化:监控每一步的进展
训练循环的代码要写得清晰且功能完整,包括训练、验证、模型保存和可视化。我习惯用TensorBoard来监控一切。
from torch.utils.tensorboard import SummaryWriter
import time
import math
def train_one_epoch(epoch, model, train_loader, criterion, optimizer, scheduler, device, writer):
model.train()
running_loss = 0.0
correct = 0
total = 0
start_time = time.time()
for batch_idx, (inputs, targets) in enumerate(train_loader):
inputs, targets = inputs.to(device), targets.to(device)
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, targets)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 统计
running_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
# 每100个batch打印一次信息,并记录到TensorBoard
if (batch_idx + 1) % 100 == 0:
batch_loss = running_loss / 100
batch_acc = 100. * correct / total
current_lr = optimizer.param_groups[0]['lr']
print(f'Epoch: {epoch} | Batch: {batch_idx+1}/{len(train_loader)} | '
f'Loss: {batch_loss:.4f} | Acc: {batch_acc:.2f}% | LR: {current_lr:.6f}')
step = epoch * len(train_loader) + batch_idx
writer.add_scalar('Train/Loss', batch_loss, step)
writer.add_scalar('Train/Accuracy', batch_acc, step)
writer.add_scalar('Train/Learning_Rate', current_lr, step)
running_loss = 0.0
scheduler.step() # 每个epoch调整一次学习率
epoch_time = time.time() - start_time
epoch_acc = 100. * correct / total
print(f'Epoch {epoch} 训练完成,耗时 {epoch_time:.0f}s,平均准确率: {epoch_acc:.2f}%')
return epoch_acc
@torch.no_grad()
def evaluate(model, test_loader, device, epoch, writer):
model.eval()
correct = 0
total = 0
start_time = time.time()
for inputs, targets in test_loader:
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
test_acc = 100. * correct / total
test_time = time.time() - start_time
print(f'测试集评估 | 准确率: {test_acc:.2f}% | 耗时: {test_time:.0f}s')
writer.add_scalar('Test/Accuracy', test_acc, epoch)
return test_acc
def main():
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'使用设备: {device}')
num_classes = len(train_dataset.idx_to_char)
model, criterion, optimizer, scheduler = create_model_and_optimizer(num_classes, device)
# TensorBoard记录器
writer = SummaryWriter(log_dir='./runs/hwdb_mobilenetv2')
best_acc = 0.0
for epoch in range(100):
print(f'\n开始第 {epoch+1}/100 轮训练')
train_acc = train_one_epoch(epoch, model, train_loader, criterion, optimizer, scheduler, device, writer)
test_acc = evaluate(model, test_loader, device, epoch, writer)
# 保存最佳模型
if test_acc > best_acc:
best_acc = test_acc
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'best_acc': best_acc,
'char_mapping': train_dataset.idx_to_char,
}, './checkpoints/best_model.pth')
print(f'* 新的最佳模型已保存,准确率: {best_acc:.2f}%')
writer.close()
print(f'训练结束,最佳测试准确率: {best_acc:.2f}%')
这个训练框架包含了完整的日志、监控和模型保存逻辑。通过TensorBoard,你可以实时查看损失和准确率曲线,分析模型是欠拟合还是过拟合,非常直观。
5. 模型优化与部署实战:从训练到上线的最后一公里
模型训练到90%以上的准确率,是不是就大功告成了?远着呢。实验室的精度只是第一步,怎么让这个模型在实际场景中又快又准地跑起来,才是真正的挑战。这部分我结合自己踩过的坑,分享几个关键的优化和部署技巧。
5.1 模型剪枝与量化:给模型“瘦身”和“加速”
MobileNetV2本身已经很小了(约3.4M参数),但在一些极端边缘设备上,我们还想让它更小、更快。剪枝和量化是两种最常用的后训练优化技术。
剪枝的原理很简单:找出网络中不重要的连接(权重接近0的),把它们去掉。PyTorch提供了相关的工具。我们可以进行简单的幅度剪枝:
import torch.nn.utils.prune as prune
def prune_model_l1_unstructured(model, pruning_rate=0.2):
"""对模型的卷积层和全连接层进行L1非结构化剪枝"""
parameters_to_prune = []
for name, module in model.named_modules():
if isinstance(module, (nn.Conv2d, nn.Linear)):
parameters_to_prune.append((module, 'weight'))
# 全局剪枝
prune.global_unstructured(
parameters_to_prune,
pruning_method=prune.L1Unstructured,
amount=pruning_rate,
)
# 重要!剪枝后需要将掩码永久化,并移除`weight_orig`和`weight_mask`
for module, _ in parameters_to_prune:
prune.remove(module, 'weight')
print(f"已完成全局非结构化剪枝,比例: {pruning_rate}")
return model
# 加载训练好的模型
checkpoint = torch.load('./checkpoints/best_model.pth', map_location='cpu')
pruned_model = MobileNetV2(num_classes=3755).to('cpu')
pruned_model.load_state_dict(checkpoint['model_state_dict'])
# 执行剪枝
pruned_model = prune_model_l1_unstructured(pruned_model, pruning_rate=0.3)
# 剪枝后需要重新微调(fine-tune)几个epoch以恢复精度
量化则是降低模型中权重和激活值的数值精度,比如从32位浮点数(FP32)降到8位整数(INT8)。这能显著减少模型大小和内存占用,并利用硬件对整型计算的支持来加速推理。PyTorch的量化支持动态量化和静态量化。对于MobileNetV2这种CNN,静态量化效果更好:
import torch.quantization
# 1. 定义量化配置
model_to_quantize = MobileNetV2(num_classes=3755)
model_to_quantize.load_state_dict(torch.load('./checkpoints/best_model.pth')['model_state_dict'])
model_to_quantize.eval()
# 2. 融合模型中的Conv+BN+ReLU模块(量化前准备)
model_to_quantize.fuse_model()
# 3. 指定量化配置
model_to_quantize.qconfig = torch.quantization.get_default_qconfig('fbgemm') # 用于x86 CPU后端
# 4. 准备量化(插入观察器,记录激活值的范围)
torch.quantization.prepare(model_to_quantize, inplace=True)
# 5. 校准(用少量数据跑一遍,让观察器记录数据分布)
# 这里需要用一些校准数据,比如测试集的一部分
calibration_data = ... # 获取一些数据
with torch.no_grad():
for data in calibration_data:
model_to_quantize(data)
# 6. 转换到量化模型
quantized_model = torch.quantization.convert(model_to_quantize, inplace=False)
# 保存量化后的模型
torch.jit.save(torch.jit.script(quantized_model), './checkpoints/quantized_mobilenetv2.pth')
print("量化模型已保存。")
量化后的模型,推理速度在支持INT8的硬件上能有数倍提升,模型文件大小也能减少约75%。不过要注意,量化可能会带来轻微的精度损失,需要仔细校准。
5.2 使用ONNX进行跨平台部署
在实际项目中,你的模型可能需要在不同的推理引擎(如TensorRT, OpenVINO, NCNN)或不同的硬件(如NVIDIA GPU, Intel CPU, ARM NPU)上运行。ONNX是一个开放的模型格式标准,可以充当这个“中间人”。
import torch.onnx
# 加载模型
model = MobileNetV2(num_classes=3755)
model.load_state_dict(torch.load('./checkpoints/best_model.pth', map_location='cpu')['model_state_dict'])
model.eval()
# 创建一个示例输入
dummy_input = torch.randn(1, 3, 32, 32)
# 导出为ONNX格式
torch.onnx.export(model,
dummy_input,
"./checkpoints/mobilenetv2_hwdb.onnx",
export_params=True,
opset_version=13, # 使用较新的算子集
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}, # 支持动态batch
'output': {0: 'batch_size'}})
print("ONNX模型导出成功。")
导出的ONNX模型,你就可以用ONNX Runtime在各种平台上进行高性能推理了。例如,用Python进行推理:
import onnxruntime as ort
import numpy as np
# 创建ONNX Runtime推理会话
ort_session = ort.InferenceSession('./checkpoints/mobilenetv2_hwdb.onnx')
# 准备输入数据(需要是numpy array)
input_name = ort_session.get_inputs()[0].name
dummy_input_np = np.random.randn(1, 3, 32, 32).astype(np.float32)
# 运行推理
outputs = ort_session.run(None, {input_name: dummy_input_np})
print(f"推理输出形状: {outputs[0].shape}")
5.3 构建一个简单的推理服务
最后,我们把这个模型包装成一个简单的Web API,模拟一个真实的部署场景。这里我们用Flask快速搭建一个服务。
from flask import Flask, request, jsonify
import onnxruntime as ort
import numpy as np
from PIL import Image
import io
app = Flask(__name__)
# 加载模型和标签映射
ort_session = ort.InferenceSession('./checkpoints/mobilenetv2_hwdb.onnx')
with open('./HWDB1.1_dataset/char_dict_train.pkl', 'rb') as f:
char_dict = pickle.load(f)
idx_to_char = char_dict['idx_to_char']
def preprocess_image(image_bytes):
"""预处理上传的图片"""
image = Image.open(io.BytesIO(image_bytes)).convert('L').convert('RGB')
# 应用与训练时相同的转换(除数据增强外)
transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
image_tensor = transform(image).unsqueeze(0) # 增加batch维度
return image_tensor.numpy()
@app.route('/predict', methods=['POST'])
def predict():
if 'file' not in request.files:
return jsonify({'error': 'No file uploaded'}), 400
file = request.files['file']
if file.filename == '':
return jsonify({'error': 'No selected file'}), 400
try:
# 1. 读取并预处理图片
img_bytes = file.read()
input_np = preprocess_image(img_bytes)
# 2. 运行模型推理
input_name = ort_session.get_inputs()[0].name
outputs = ort_session.run(None, {input_name: input_np})
predictions = outputs[0][0] # 取第一个batch的结果
# 3. 获取Top-5预测结果
top5_idx = np.argsort(predictions)[-5:][::-1]
top5_conf = predictions[top5_idx]
top5_chars = [idx_to_char[idx] for idx in top5_idx]
result = {
'predictions': [
{'character': char, 'confidence': float(conf)}
for char, conf in zip(top5_chars, top5_conf)
]
}
return jsonify(result)
except Exception as e:
return jsonify({'error': str(e)}), 500
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000, debug=False)
这个简单的服务接收一个图片文件,返回识别出的汉字及其置信度。你可以用Postman或者写个简单的Python脚本来测试它。这只是一个起点,在生产环境中,你还需要考虑模型版本管理、请求队列、负载均衡、监控告警等一系列问题。
走到这一步,一个完整的手写汉字识别系统,从数据准备、模型构建、训练优化到部署上线的全链路,你就都走通了。这个过程里,每一个环节都有坑,也都有优化的空间。我自己的体会是,深度学习项目,代码跑通只是第一步,如何让它在真实场景下稳定、高效、可靠地运行,才是更考验工程能力的地方。希望我分享的这些实战经验和代码片段,能帮你少走些弯路。如果在复现过程中遇到什么问题,欢迎随时交流。
更多推荐
所有评论(0)