pytorch入门实战之验证码识别

本文将通过pytorch框架训练一个四层卷积神经网络&＃xff0c;用以识别四位数字字母区分大小的验证码。使用800张验证码图片做为训练集&＃xff0c;准确率最高达75%。译

1. 引言 去年四六级查分时候我把准考证号忘了&＃xff0c;准考证一时也找不到&＃xff0c;最后是靠试准考证号试出来的&＃xff0c;因为和我同一个考场的同学准考证号只有最后两位座位号不一样&＃xff0c;一个考场不超过30人&＃xff0c;遍历座位号就能试出来。

   四六级查分系统有一个四位数字字母验证码&＃xff0c;如果能够自动识别验证码&＃xff0c;就能不断遍历准考证号查分了&＃xff0c;不用手动输入验证码查分&＃xff0c;效率大大提高&＃xff0c;不知道淘宝上“忘记准考证号帮查四六级分”服务是不是这样做的。

2. 数据收集 四六级查分网页链接为http://cet.neea.edu.cn/cet&＃xff0c;查分网页如下图&＃xff1a;

   首先按Fn&＃43;F12使用网页开发者工具抓包看一下验证码是如何请求&＃xff0c;以及如何提交查询信息并返回结果。最好不要一次性把三条信息都输对&＃xff0c;那样会直接跳到查询结果页&＃xff0c;不方便查看提交查询的请求。

   可以很容易的找到提交请求的是一个post请求&＃xff0c;请求地址为http://cache.neea.edu.cn/cet/query&＃xff0c;请求参数有两个&＃xff0c;分别是data和v&＃xff0c;data是由一串固定字符和准考证号以及姓名组成&＃xff0c;v则是验证码。通过构建查询请求&＃xff0c;我们可以知道验证码是否输入正确。点击获取验证码按钮&＃xff0c;可以抓包获取到验证码的请求&＃xff0c;将验证码请求以及提交查询写成python代码如下&＃xff1a;

[代码]py代码&＃xff1a; def get_captcha_img(): ik &＃61; &＃39;123456789123456&＃39; rand &＃61; random.random() img_path &＃61; &＃39;{}/{}.png&＃39;.format(false_dir, rand) imgs_url &＃61; &＃39;http://cache.neea.edu.cn/Imgs.do?c&＃61;CET&ik&＃61;{}&t&＃61;{}&＃39;.format(ik, rand) headers &＃61; {&＃39;Referer&＃39;: &＃39;http://cet.neea.edu.cn/cet&＃39;} resp &＃61; sess.get(imgs_url, headers&＃61;headers) img_url &＃61; re.findall(r&＃39;"([^"]*)?"&＃39;, resp.text)[0] img_resp &＃61; sess.get(img_url, headers&＃61;headers) with open(img_path, &＃39;wb&＃39;) as f: f.write(img_resp.content) return img_path

def check_captcha(v): query_url &＃61; &＃39;http://cache.neea.edu.cn/cet/query&＃39; data &＃61; {&＃39;data&＃39;: &＃39;CET4_181_DANGCI,123456789123456,萧炎&＃39;, &＃39;v&＃39;: v} headers &＃61; {&＃39;Referer&＃39;: &＃39;http://cet.neea.edu.cn/cet&＃39;} resp &＃61; sess.post(query_url, headers&＃61;headers, data&＃61;data) if &＃39;抱歉&＃xff0c;验证码错误&＃xff01;&＃39; in resp.text: return False else: return True

结合以上请求验证码以及提交查询信息判断验证码是否正确的方法&＃xff0c;再通过打码平台&＃xff0c;可以获得带有正确标记的验证码图片。 使用上述方法&＃xff0c;我获得了1181张带有标注的验证码&＃xff0c;宽和高为&＃xff08;180&＃xff0c;100&＃xff09;&＃xff0c;将其分为训练集与测试集&＃xff0c;训练集为800张&＃xff0c;测试及381张。我看过的**使用卷积神经网络识别验证码的文章&＃xff0c;使用的训练集数量多达几千上万张&＃xff0c;大多都是自己用程序生成的&＃xff0c;本文使用打码平台标记的验证码&＃xff0c;就不要求那么大的数据集了&＃xff0c;但也能达到满意的效果。还值得一提的是&＃xff0c;使用打码平台标注验证码&＃xff0c;成功标注了1181张外&＃xff0c;还有将近四百张验证码识别失败&＃xff0c;粗略估计&＃xff0c;这个打码平台准确率在75%左右。


3. CNN模型搭建 CNN主要由卷积层&＃xff0c;池化层&＃xff0c;激活函数组成&＃xff0c;再加上一个BatchNorm&＃xff0c;BatchNorm叫做批规范化&＃xff0c;可以加速模型的收敛速度。

   模型代码如下&＃xff1a;

[代码]py代码&＃xff1a; import torch.nn as nn

class CNN(nn.Module): def init(self, num_class&＃61;36, num_char&＃61;4): super(CNN, self).init() self.num_class &＃61; num_class self.num_char &＃61; num_char self.conv &＃61; nn.Sequential( #batch3180100 nn.Conv2d(3, 16, 3, padding&＃61;(1, 1)), nn.MaxPool2d(2, 2), nn.BatchNorm2d(16), nn.ReLU(), #batch169050 nn.Conv2d(16, 64, 3, padding&＃61;(1, 1)), nn.MaxPool2d(2, 2), nn.BatchNorm2d(64), nn.ReLU(), #batch644525 nn.Conv2d(64, 512, 3, padding&＃61;(1, 1)), nn.MaxPool2d(2, 2), nn.BatchNorm2d(512), nn.ReLU(), #batch5122212 nn.Conv2d(512, 512, 3, padding&＃61;(1, 1)), nn.MaxPool2d(2, 2), nn.BatchNorm2d(512), nn.ReLU(), #batch512116 ) self.fc &＃61; nn.Linear(512116, self.num_classself.num_char)

def forward(self, x):x &＃61; self.conv(x)x &＃61; x.view(-1, 512*11*6)x &＃61; self.fc(x)return x
nn.Sequential()可以看作模块的有序容器&＃xff0c;可以方便快捷的搭建神经网络。


网络的输入是一个shape为[batch, 3, 180, 100]的张量&＃xff0c;batch代表的是一个批次图片数量&＃xff0c;3代表输入的图片是3通道的&＃xff0c;即RGB&＃xff0c;180和100则分别代表图片的宽和高。

主要的结构如下&＃xff1a;


第一个卷积层nn.Conv2d(3, 16, 3, padding&＃61;(1, 1))&＃xff0c;参数分别对应着输入的通道数3&＃xff0c;输出通道数16&＃xff0c;卷积核大小为3&＃xff08;长宽都为3&＃xff09;&＃xff0c;padding为&＃xff08;1&＃xff0c; 1&＃xff09;可以保证输入输出的长宽不变。shape为[batch, 3, 180, 100]的张量通过这个卷积层&＃xff0c;输出一个shape为[batch, 16, 180, 100]的张量。

接着一个最大池化层nn.MaxPool2d(2, 2)&＃xff0c;参数分别对应着池化窗口大小为2&＃xff08;长宽都为2&＃xff09;&＃xff0c;步长为3. 输出的长宽为输入的一半&＃xff0c;如果长宽为奇数的话则补边。输入一个shape为[batch, 16, 180, 100]的张量&＃xff0c;输出为一个shape为[batch, 16, 90, 50]的张量。

批规范层nn.BatchNorm2d(16)&＃xff0c;16为输入张量的通道数。 激活函数nn.ReLu()&＃xff0c;就是把小于0的值置0&＃xff0c;大于0的值不变&＃xff0c;使用激活函数是为了引入非线性&＃xff0c;让模型可以拟合更复杂的函数。

经过4组如上结构的卷积后&＃xff0c;得到一个shape为[batch, 512, 11, 6]的张量&＃xff0c;x.view(-1, 512*11*6)将改变张量的shape为[batch, 512*11*6]&＃xff0c;再用一个[512*11*6, num_class*num_char]的全连接层映射为一个[batch, num_class*num_char]张量&＃xff0c;这个就是模型的输出&＃xff0c;其中num_class代表字符的种类数量&＃xff0c;num_char代表一张验证码图片含有的字符数量&＃xff0c;分别为36与4。


4. 数据加载 pytorch有非常方便高效的数据加载模块--Dataset和DataLoader。 Dataset是数据样本的封装&＃xff0c;可以很方便的读取数据。实现一个Dataset的子类&＃xff0c;需要重写__len__和__getitem__方法&＃xff0c;__len__需要返回整个数据集的大小&＃xff0c;__getitem__提供一个整数索引参数&＃xff0c;一个样本数据&＃xff08;一个图片张量和一个标签张量&＃xff09;。 验证码图片的Dataset代码如下&＃xff1a;

[代码]py代码&＃xff1a; class CaptchaData(Dataset): def init(self, data_path, num_class&＃61;36, num_char&＃61;4, transform&＃61;None, target_transform&＃61;None, alphabet&＃61;alphabet): super(Dataset, self).init() self.data_path &＃61; data_path self.num_class &＃61; num_class self.num_char &＃61; num_char self.transform &＃61; transform self.target_transform &＃61; target_transform self.alphabet &＃61; alphabet self.samples &＃61; 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 &＃61; self.samples[index]img &＃61; img_loader(img_path)if self.transform is not None:img &＃61; self.transform(img)if self.target_transform is not None:target &＃61; self.target_transform(target)return img, torch.Tensor(target)
其中make_dataset为读取图片路径和标签的函数&＃xff0c;返回[(img_path, target), (img_path, target), ...]的数据形式。img_loader为读取图片的函数&＃xff0c;并且转换成RGB三通道。 这两个函数具体实现如下&＃xff1a;


[代码]py代码&＃xff1a; def img_loader(img_path): img &＃61; Image.open(img_path) return img.convert(&＃39;RGB&＃39;)

def make_dataset(data_path, alphabet, num_class, num_char): img_names &＃61; os.listdir(data_path) samples &＃61; [] for img_name in img_names: img_path &＃61; os.path.join(data_path, img_name) target_str &＃61; img_name.split(&＃39;.&＃39;)[0] assert len(target_str) &＃61;&＃61; num_char target &＃61; [] for char in target_str: vec &＃61; [0] * num_class vec[alphabet.find(char)] &＃61; 1 target &＃43;&＃61; vec samples.append((img_path, target)) return samples DataLoader是Dataset的进一步封装&＃xff0c;Dataset每次通过__getitem__方法取到的是一个样本&＃xff0c;经过DataLoader封装为dataloader后&＃xff0c;每次取的是一个batch大小的样本批次。

[代码]py代码&＃xff1a; transforms &＃61; Compose([ToTensor()]) train_dataset &＃61; CaptchaData(&＃39;./data/train&＃39;, transform&＃61;transforms) train_data_loader &＃61; DataLoader(train_dataset, batch_size&＃61;batch_size, num_workers&＃61;0, shuffle&＃61;True, drop_last&＃61;True) test_data &＃61; CaptchaData(&＃39;./data/test&＃39;, transform&＃61;transforms) test_data_loader &＃61; DataLoader(test_data, batch_size&＃61;batch_size, num_workers&＃61;0, shuffle&＃61;True, drop_last&＃61;True) transform是数据预处理操作&＃xff0c;一般数据增强就通过transform实现&＃xff0c;可以随机亮度&＃xff0c;随机翻转&＃xff0c;随机缩放等等。此处只使用了ToTensor()&＃xff0c;将PIL.Image对象转换成Tensor。

5. 训练 训练网络的一般流程为&＃xff1a;
   1. 定义网络
   2. 定义优化器optimizer和损失函数criterion
   3. 遍历dataloader&＃xff0c;每次取一个batch训练。计算loss&＃xff0c;将优化器梯度置零&＃xff0c;loss向后传播&＃xff0c;计算梯度&＃xff0c;优化器更新参数。
   4. 训练集训练完一个epoch后&＃xff0c;使用测试集计算下准确率。
   5. 保存模型 主要代码如下&＃xff1a; [代码]py代码&＃xff1a; cnn &＃61; CNN() if torch.cuda.is_available(): cnn.cuda() optimizer &＃61; torch.optim.Adam(cnn.parameters(), lr&＃61;base_lr) criterion &＃61; nn.MultiLabelSoftMarginLoss()

for epoch in range(max_epoch): cnn.train() for img, target in train_data_loader: img &＃61; Variable(img) target &＃61; Variable(target) if torch.cuda.is_available(): img &＃61; img.cuda() target &＃61; target.cuda() output &＃61; cnn(img) loss &＃61; criterion(output, target) optimizer.zero_grad() loss.backward() optimizer.step()

loss_history &＃61; []
acc_history &＃61; []
cnn.eval()
for img, target in test_data_loader:img &＃61; Variable(img)target &＃61; Variable(target)if torch.cuda.is_available():img &＃61; img.cuda()target &＃61; target.cuda()output &＃61; cnn(img)acc &＃61; calculat_acc(output, target)acc_history.append(acc)loss_history.append(float(loss))
torch.save(cnn.state_dict(), model_path)
其中&＃xff0c;cnn.train()将网络切换到训练状态&＃xff0c;cnn.eval()将网络切换到模型评估状态&＃xff0c;这两者的差别主要体现在dropout和batchnorm层中&＃xff0c;模型评估状态下&＃xff0c;将不会启用dropout层&＃xff0c;batchnrom不会更新均值和标准差。cnn.cuda()将数据张量分配到cuda设备上&＃xff08;英伟达显卡&＃xff09;&＃xff0c;加快运算速度。 损失函数使用的是nn.MultiLabelSoftMarginLoss()&＃xff0c;多分类多标签损失函数。每个类别有多个标签&＃xff0c;集每张验证码有四个字符。选择accuracy&＃xff08;预测准确率&＃xff09;做为模型的评估指标&＃xff0c;需要再编写一个计算准确率的函数&＃xff1a;


[代码]py代码&＃xff1a; def calculat_acc(output, target): output, target &＃61; output.view(-1, 36), target.view(-1, 36) output &＃61; nn.functional.softmax(output, dim&＃61;1) output &＃61; torch.argmax(output, dim&＃61;1) target &＃61; torch.argmax(target, dim&＃61;1) output, target &＃61; output.view(-1, 4), target.view(-1, 4) correct_list &＃61; [] for i, j in zip(target, output): if torch.equal(i, j): correct_list.append(1) else: correct_list.append(0) acc &＃61; sum(correct_list) / len(correct_list) return acc
训练结果&＃xff1a;

最终训练了五十几个epoch后&＃xff0c;测试集准确率最高达75%&＃xff0c;训练集已过拟合达100%。再将验证码打印出来&＃xff0c;预测与实际标签对比&＃xff1a;


6. 结语 仅使用800张验证码图片做为训练集&＃xff0c;就能最终达到75%的准确率&＃xff0c;效果还是比较满意的&＃xff0c;已经和打码平台差不多了。要想进一步的提高准确率&＃xff0c;需要扩充数据集。可以将已经训练好&＃xff0c;准确率达到75%的模型代替打码平台&＃xff0c;去获取更多标注好的验证码。数据集充分的情况下&＃xff0c;准确率达到90%是比较容易的。

   完整代码可以从Github上获取&＃xff1a;https://github.com/ice-tong/pytorch-captcha

   数据集可以从百度网盘下载&＃xff1a; 链接: https://pan.baidu.com/s/1pHSl-5nHJWazXVqda-2IcA 提取码: mv3u

   文&＃xff1a;icetong 更多人工智能相关文章&＃xff1a;http://siligence.ai