pytorch⼊门实战之验证码识别
本⽂将通过pytorch框架训练⼀个四层卷积神经⽹络,⽤以识别四位数字字母区分⼤⼩的验证码。使⽤800张验证码图⽚做为训练集,准确
率最⾼达75%。译
1. 引⾔ 去年四六级查分时候我把准考证号忘了,准考证⼀时也不到,最后是靠试准考证号试出来的,因为和我同⼀个考场的同学准考
证号只有最后两位座位号不⼀样,⼀个考场不超过30⼈,遍历座位号就能试出来。
四六级查分系统有⼀个四位数字字母验证码,如果能够⾃动识别验证码,就能不断遍历准考证号查分了,不⽤⼿动输⼊验证码查分,
效率⼤⼤提⾼,不知道淘宝上“忘记准考证号帮查四六级分”服务是不是这样做的。
2. ⾸先按Fn+F12使⽤⽹页开发者⼯具抓包看⼀下验证码是如何请求,以及如何提交查询信息并返回结果。最好不要⼀次性把三条信息
都输对,那样会直接跳到查询结果页,不⽅便查看提交查询的请求。
结合以上请求验证码以及提交查询信息判断验证码是否正确的⽅法,再通过打码平台,可以获得带有正确标记的验证码图⽚。使⽤上述⽅法,我获得了1181张带有标注还值得⼀提的是,使⽤打码平台标注验证码,成功标注了1181张外,还有将近四百张验证码识别失败,粗略估计,这个打码平台准确率在75%左右。
3. CNN模型搭建 CNN主要由卷积层,池化层,激活函数组成,再加上⼀个BatchNorm,BatchNorm叫做批规范化,可以加速模型的
收敛速度。
模型代码如下:
[代码]py代码: as nn
class CNN(nn.Module): def init(self, num_class=36, num_char=4): super(CNN, self).init() self.num_class = num_class
self.num_char = num_v = nn.Sequential( #batch3180100 nn.Conv2d(3, 16, 3, padding=(1, 1)),
nn.MaxPool2d(2, 2), nn.BatchNorm2d(16), nn.ReLU(), #batch169050 nn.Conv2d(16, 64, 3, padding=(1, 1)),
nn.MaxPool2d(2, 2), nn.BatchNorm2d(64), nn.ReLU(), #batch644525 nn.Conv2d(64, 512, 3, padding=(1, 1)),
nn.MaxPool2d(2, 2), nn.BatchNorm2d(512), nn.ReLU(), #batch5122212 nn.Conv2d(512, 512, 3, padding=(1, 1)),
nn.MaxPool2d(2, 2), nn.BatchNorm2d(512), nn.ReLU(), #batch512116 ) self.fc = nn.Linear(512116,
self.num_class self.num_char)
def forward(self, x):
x = v(x)
x = x.view(-1, 512*11*6)
x = self.fc(x)
return x
nn.Sequential()可以看作模块的有序容器,可以⽅便快捷的搭建神经⽹络。
⽹络的输⼊是⼀个shape为[batch, 3, 180, 100]的张量,batch代表的是⼀个批次图⽚数量,3代表输⼊的图⽚是3通道的,即
RGB,180和100则分别代表图⽚的宽和⾼。
主要的结构如下:
第⼀个卷积层nn.Conv2d(3, 16, 3, padding=(1, 1)),参数分别对应着输⼊的通道数3,输出通道数16,卷积核⼤⼩为3(长宽都为
3),padding为(1, 1)可以保证输⼊输出的长宽不变。shape为[batch, 3, 180, 100]的张量通过这个卷积层,输出⼀个shape为
[batch, 16, 180, 100]的张量。
接着⼀个最⼤池化层nn.MaxPool2d(2, 2),参数分别对应着池化窗⼝⼤⼩为2(长宽都为2),步长为3. 输出的长宽为输⼊的⼀半,如果
长宽为奇数的话则补边。输⼊⼀个shape为[batch, 16, 180, 100]的张量,输出为⼀个shape为[batch, 16, 90, 50]的张量。
批规范层nn.BatchNorm2d(16),16为输⼊张量的通道数。 激活函数nn.ReLu(),就是把⼩于0的值置0,⼤于0的值不变,使⽤激活函数
是为了引⼊⾮线性,让模型可以拟合更复杂的函数。
经过4组如上结构的卷积后,得到⼀个shape为[batch, 512, 11, 6]的张量,x.view(-1, 512*11*6)将改变张量的shape为[batch, 512*11*6],再⽤⼀个[512*11*6, num_cla
4. 数据加载 pytorch有⾮常⽅便⾼效的数据加载模块--Dataset和DataLoader。 Dataset是数据样本的封装,可以很⽅便的读取数据。
实现⼀个Dataset的⼦类,需要重写__len__和__getitem__⽅法,__len__需要返回整个数据集的⼤⼩,__getitem__提供⼀个整数索
引参数,⼀个样本数据(⼀个图⽚张量和⼀个标签张量)。 验证码图⽚的Dataset代码如下:
[代码]py代码: class CaptchaData(Dataset): def init(self, data_path, num_class=36, num_char=4, transform=None,
target_transform=None, alphabet=alphabet): super(Dataset, self).init() self.data_path = data_path self.num_class =
num_class self.num_char = num_ansform = transform self.target_transform = target_transform self.alphabet =
alphabet self.samples = make_dataset(self.data_path, self.alphabet, self.num_class, self.num_char)
def __len__(self):
return len(self.samples)
def __getitem__(self, index):
img_path, target = self.samples[index]
img = img_loader(img_path)
ansform is not None:
img = ansform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, torch.Tensor(target)
其中make_dataset为读取图⽚路径和标签的函数,返回[(img_path, target), (img_path, target), ...]的数据形式。img_loader为读取图⽚的函数,并且转换成RGB三通道
[代码]py代码: def img_loader(img_path): img = Image.open(img_path) vert('RGB')
def make_dataset(data_path, alphabet, num_class, num_char): img_names = os.listdir(data_path) samples = [] for img_name
in img_names: img_path = os.path.join(data_path, img_name) target_str = img_name.split('.')[0] assert len(target_str) ==
num_char target = [] for char in target_str: vec = [0] * num_class vec[alphabet.find(char)] = 1 target += vec
samples.append((img_path, target)) return samples DataLoader是Dataset的进⼀步封装,Dataset每次通过__getitem__⽅法取到
的是⼀个样本,经过DataLoader封装为dataloader后,每次取的是⼀个batch⼤⼩的样本批次。
[代码]py代码: transforms = Compose([ToTensor()]) train_dataset = CaptchaData('./data/train', transform=transforms)
train_data_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=0, shuffle=True, drop_last=True)
test_data = CaptchaData('./data/test', transform=transforms) test_data_loader = DataLoader(test_data,
batch_size=batch_size, num_workers=0, shuffle=True, drop_last=True) transform是数据预处理操作,⼀般数据增强就通过
transform实现,可以随机亮度,随机翻转,随机缩放等等。此处只使⽤了ToTensor(),将PIL.Image对象转换成Tensor。
5. 训练 训练⽹络的⼀般流程为:
1. 定义⽹络
2. 定义优化器optimizer和损失函数criterion
3. 遍历dataloader,每次取⼀个batch训练。计算loss,将优化器梯度置零,loss向后传播,计算梯度,优化器更新参数。
4. 训练集训练完⼀个epoch后,使⽤测试集计算下准确率。
5. 保存模型 主要代码如下: [代码]py代码: cnn = CNN() if torch.cuda.is_available(): cnn.cuda() optimizer =
torch.optim.Adam(cnn.parameters(), lr=base_lr) criterion = nn.MultiLabelSoftMarginLoss()
for epoch in range(max_epoch): ain() for img, target in train_data_loader: img = Variable(img) target = Variable(target)
if torch.cuda.is_available(): img = img.cuda() target = target.cuda() output = cnn(img) loss = criterion(output, target)
<_grad() loss.backward() optimizer.step()
loss_history = []
acc_history = []
cnn.eval()
for img, target in test_data_loader:
img = Variable(img)
target = Variable(target)
if torch.cuda.is_available():
img = img.cuda()
target = target.cuda()
准考证号忘了怎么查询
output = cnn(img)
acc = calculat_acc(output, target)
acc_history.append(acc)
loss_history.append(float(loss))
torch.save(cnn.state_dict(), model_path)
其中,ain()将⽹络切换到训练状态,cnn.eval()将⽹络切换到模型评估状态,这两者的差别主要体现在dropout和batchnorm层中,模型评估状态下,将不会启⽤d 选择accuracy(预测准确率)做为模型的评估指标,需要再编写⼀个计算准确率的函数:
[代码]py代码: def calculat_acc(output, target): output, target = output.view(-1, 36), target.view(-1, 36) output =
nn.functional.softmax(output, dim=1) output = torch.argmax(output, dim=1) target = torch.argmax(target, dim=1) output,
target = output.view(-1, 4), target.view(-1, 4) correct_list = [] for i, j in zip(target, output): if torch.equal(i, j):
correct_list.append(1) else: correct_list.append(0) acc = sum(correct_list) / len(correct_list) return acc
训练结果:
最终训练了五⼗⼏个epoch后,测试集准确率最⾼达75%,训练集已过拟合达100%。
再将验证码打印出来,预测与实际标签对⽐:
6. 结语 仅使⽤800张验证码图⽚做为训练集,就能最终达到75%的准确率,效果还是⽐较满意的,已经和打码平台差不多了。要想进⼀
步的提⾼准确率,需要扩充数据集。可以将已经训练好,准确率达到75%的模型代替打码平台,去获取更多标注好的验证码。数据集
充分的情况下,准确率达到90%是⽐较容易的。