前言
大家好，我是阿光。
本专栏整理了《深度学习100例》，内包含了各种不同的深度学习项目，包含项目原理以及源码，每一个项目实例都附带有完整的代码+数据集。
正在更新中~ ✨
🚨 我的项目环境：
- 平台：Windows10
- 语言环境：python3.7
- 编译器：PyCharm
- PyTorch版本：1.8.1
💥 项目专栏：【PyTorch深度学习项目实战100例】
本文主要实现了两个工作：
- 1.验证码生成
- 2.PyTorch识别验证码
一、生成验证码图片
我们利用开源库captcha进行生成验证码
1.随机生成数字
首先自定义一个函数进行随机生成验证码，这里只是生成4个数字，可以自己加一些字母或者符号。
def random_captcha_text(char_set=number, captcha_size=4): # 可以设置只用来生成数字
captcha_text = []
for i in range(captcha_size):
c = random.choice(char_set)
captcha_text.append(c)
return captcha_text
2.将标签以及验证码保存
生成好随机数字之后，然后利用captcha生成对应的验证码，然后将这个验证码保存到指定目录，然后再将验证码保存到定义的标签文件中。
def gen_capthcha_text_and_image(m):
image = ImageCaptcha()
captcha_text = random_captcha_text() # 生成数字
captcha_text = ' '.join(captcha_text) # 生成标签
image.write(captcha_text, "./data/image/" + '%.4d' % m + '.jpg') # 保存图片
# 将标签信息写入
with open(data_path + "label.txt", "a") as f:
f.write(captcha_text)
f.writelines("\n")
在这里插入图片描述
二、PyTorch进行识别
对于一张验证码来讲，上面有4个数字，所以我们可以将每个图片对应的输出为40维，由于是识别数字，可以将每个数字的识别看成10分类，利用交叉熵来实现，所以每个图片就是40维度。
1.自定义数据集
这里我们需要重写DataSet类，加载我们的验证码数据和label标签文件。
2.自定义网络
我们自已利用PyTorch搭建了简单的卷积网络，由于验证码比较简单，所以不需要特殊的网络模型。 首先需要几个卷积进行特征提取，然后在尾部加入全连接层，最后一层的输出维度为40维度。
# 自定义卷积网络模型
class ConvNet(nn.Module):
def __init__(self):
super(ConvNet, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=4, stride=1, padding=2), # 验证码的大小为(3,60,160)
nn.BatchNorm2d(32),
nn.LeakyReLU(0.2, inplace=True),
nn.MaxPool2d(kernel_size=2), # out:(bs,32,30,80)
nn.Conv2d(32, 64, kernel_size=4, stride=1, padding=2),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2, inplace=True),
nn.MaxPool2d(kernel_size=2), # out:(bs,64,15,40)
nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2, inplace=True),
nn.MaxPool2d(kernel_size=2) # out:(bs,64,7,20)
)
self.fc1 = nn.Linear(64 * 7 * 20, 500)
self.fc2 = nn.Linear(500, 40) # 每个图片中有4个数字，每个数字为10分类，所以为40个输出
def forward(self, x):
x = self.conv(x)
x = x.view(x.size(0), -1) # reshape to (batch_size, 64 * 7 * 30)
output = self.fc1(x)
output = self.fc2(output)
return output
3.自定义损失函数
由于每个图像中含有4个数字，且每个数字对应10维，所以需要自己实现损失函数来完成我们的任务要求。 首先在output中提取出第一个数字对应的10个维度，然后将后面每个数字对应的10个维度按照行进行堆叠，形成4个样本数据，然后计算其损失，利用交叉熵函数，label同理。
4.模型训练
net = ConvNet()
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
loss_func = nCrossEntropyLoss()
# 最好模型的权重，以及准确率
best_model_wts = net.state_dict()
best_acc = 0.0
since = time.time()
for epoch in range(EPOCH):
running_loss = 0.0
running_corrects = 0
for step, (inputs, label) in enumerate(dataloader):
pred = torch.LongTensor(BATCH_SIZE, 1).zero_()
inputs = Variable(inputs) # (bs, 3, 60, 240)
label = Variable(label) # (bs, 4)
optimizer.zero_grad()
output = net(inputs) # (bs, 40)
loss = loss_func(output, label)
for i in range(4):
pre = F.log_softmax(output[:, 10 * i:10 * i + 10], dim=1) # (bs, 10)
pred = torch.cat((pred, pre.data.max(1, keepdim=True)[1].cpu()), dim=1) #
loss.backward()
optimizer.step()
running_loss += loss.data * inputs.size()[0]
# running_loss += loss.data[0] * inputs.size()[0]
running_corrects += equal(pred.numpy()[:, 1:], label.data.cpu().numpy().astype(int))
epoch_loss = running_loss / dataset_size
epoch_acc = running_corrects / dataset_size
if epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = net.state_dict()
if epoch == EPOCH - 1:
# 保存模型
torch.save(best_model_wts, file_path + '/best_model.pkl')
print()
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Train Loss:{:.4f} Acc: {:.4f}'.format(epoch_loss, epoch_acc))
三、验证码识别
首先读取我们需要识别的验证码，由于我们的验证码的维度为[3,60,160]，所以我们需要使用unsqueeze进行维度升维，第一维变成批次[1,3,60,160]。
img_path = r'E:\Code\Jupyter\《深度学习100例》\4.pytorch识别验证码\data\image\0000.jpg'
img = Image.open(img_path)
img = transforms.ToTensor()(img)
img = torch.unsqueeze(img, dim=0)
在这里插入图片描述
# 网络结构
class ConvNet(nn.Module):
def __init__(self):
super(ConvNet, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=4, stride=1, padding=2), # in:(bs,3,60,160)
nn.BatchNorm2d(32),
nn.LeakyReLU(0.2, inplace=True),
nn.MaxPool2d(kernel_size=2), # out:(bs,32,30,80)
nn.Conv2d(32, 64, kernel_size=4, stride=1, padding=2),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2, inplace=True),
nn.MaxPool2d(kernel_size=2), # out:(bs,64,15,40)
nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2, inplace=True),
nn.MaxPool2d(kernel_size=2) # out:(bs,64,7,20)
)
self.fc1 = nn.Linear(64 * 7 * 20, 500)
self.fc2 = nn.Linear(500, 40)
def forward(self, x):
x = self.conv(x)
x = x.view(x.size(0), -1) # reshape to (batch_size, 64 * 7 * 30)
output = self.fc1(x)
output = self.fc2(output)
return output
加载我们训练好的模型，然后用于进行识别。
# 加载模型权重
model = ConvNet()
model.load_state_dict(torch.load('./data/image/best_model.pkl'))
将一张图片输出的40维拆分成4个数字对应的10维，也就是[4,10]，然后每个数字对应的样本的最大值所在的索引提出就是对应的数字。
完整源码
【PyTorch深度学习项目实战100例】—— 使用PyTorch实现验证码识别 | 第4例_咕 嘟的博客-CSDN博客
