CNN(一)(终章)——LeNet实战(英文+数字识别)

LeNet pytorch实战续集

编写train训练模块

万事俱备，只欠训练。效果到底好不好，试一试才知道。

话不多说，上代码：

def train(epochs):  # 定义每个epoch的训练细节。
    model.train()  # 设置为training模式。
    train_acc = 0  # 初始化精确度。
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = Variable(data), Variable(target)  # 把数据转换成Variable。
        optimizer.zero_grad()  # 优化器梯度初始化为零。
        output = model(data)  # 把数据输入网络并得到输出，即进行前向传播。
        loss = loss_function(output, target)  # 交叉熵损失函数。
        loss.backward()  # 反向传播梯度。
        optimizer.step()  # 结束一次前传+反传之后，更新参数。
        pred = torch.max(output, 1)[1]
        train_acc += pred.eq(target.data.view_as(pred)).cpu().sum()  # 精确度叠加。
        if batch_idx % 10 == 0:  # 训练完成，准备打印相关信息。
            # 输出epoch，训练进度，loss，精确度。
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}, Accuracy: ({:.0f}%)\n'.format(
                epochs, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item(), 100. * train_acc / len(train_loader.dataset)))

train代码，也没啥，就是写一个循环，爬每一个batch的loss进行小批量随机梯度下降咯。。。
降降降，拟合拟合，成功成功。
再写个测试代码，一起运动。。。。

编写test测试模块

def test():
    model.eval()  # 设置为test模式。
    test_loss = 0  # 初始化测试损失值为0。
    correct = 0  # 初始化预测正确的数据个数为0。
    for data, target in test_loader:
        data, target = Variable(data), Variable(target)  # 计算前要把变量变成Variable形式，因为这样子才有梯度。

        output = model(data)  # 前向传播。
        test_loss += loss_function(output, target).item()  # sum up batch loss 把所有loss值进行累加。
        pred = output.data.max(1, keepdim=True)[1]  # get the index of the max log-probability
        correct += pred.eq(target.data.view_as(pred)).cpu().sum()  # 对预测正确的数据个数进行累加。

    test_loss /= len(test_loader.dataset)  # 因为把所有loss值进行过累加，所以最后要除以总得数据长度才得平均loss
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

测试代码也搞定了，那我们就写个入口，开始训练和测试吧~~~

运行训练和测试

# 训练：
for epoch in range(1, EPOCHS + 1):  # 以epoch为单位进行循环
    train(epoch)
# 测试：
test()
# 保存模型：
torch.save(model, 'model.pth')  # 保存模型
# 保存模型
torch.save(model.state_dict(), "model_epoch_50.pth")  # 保存权重参数

之后就开始等待吧~，

等他循环，循环，循环，

训练，训练，训练

因为我数据集不大，用不用gpu都无所谓咯，半个小时给它搞定了就是了。。我们就等着权重文件生成吧，乘着生成权重文件，写一个预测的demo吧。。。

编写预测模块

def preproccess(data):  # 定义一个数据类型转换子类。
    toTensor = transforms.ToTensor()
    data = toTensor(data)
    return data


def demo(pre_img):  # 预测任务
    checkpoint = "model_epoch_50.pth"   # 权重加载
    decode = open("decode.txt")  # 解码
    decode = list(decode)
    checkpoint = torch.load(checkpoint, map_location='cpu')  # 加载checkpoint
    model.load_state_dict(checkpoint)  # 权重输入模型
    model.eval()  # 测试模式
    tensor = preproccess(pre_img)  # 将图片转化成tensor
    tensor = Variable(torch.unsqueeze(tensor, dim=0).float(), requires_grad=False)  # 转换格式
    predict = model(tensor)[0].detach().numpy()  # 预测
    label = np.argmax(predict[:], axis=0)
    label = decode[int(label)]   # 输出标签
    return label

注意看demo第二行，用到了一个简单的解码本。解码本主要是将10~36字母部分进行一个解码而已咯。。
解码本内容很简单，部分如下：

那我又想 识别一个图片，又想识别一个文件夹呢，然后拼接呢，

码，继续码。

 # 识别文件夹下文件
    image_path = "images"  # 输入图片路径
    img_path = os.listdir(image_path)
    img_path.sort()
    img_path.sort(key=lambda x: int(x[:-4]))
    labels = []
    label1 = []
    for i in img_path:   # 遍历路径
        label = demo(Image.open(image_path + '/' + i))  # 找到图片输送到预测demo，生成标签
        print("picture:{}, label:{}".format(i, label))
        labels.append(label)
    for line in labels:
        word = line.split()[0]
        label1.append(word)
    labelx = label1[0]
    for i in label1[1:]:
        labelx += i
    print("labels:{}".format(labelx))

    # 识别单个图片
    image_path = "123.png"
    single_image = Image.open(image_path)
    label_single_image = demo(single_image)
    print("image:{},label:{}".format(image_path, label_single_image))

这些我就不做备注了，按我这python水平，写的都是bug，能用就行。噗~

差不多了，代码写完，权重文件也来了。开始预测一下吧！

看了一下训练的结果，

不错不错，测试精度95%。开始预测。

预测

放入文件夹下图片

预测结果：

结果看起来不错啊，达到我的目的了

结束结束，LeNet在字符识别上，还是不错的~~