卷积神经网络——LeNet（pytorch实现）

前言

LeNet，它是最早发布的卷积神经网络之一。LeNet被广泛用于自动取款机（ATM）机中，帮助识别处理支票的数字。 时至今日，一些自动取款机仍在运行Yann LeCun和他的同事Leon Bottou在上世纪90年代写的代码。本文在CIFAR10上训练和测试，我记得最初的LeNet是在MNIST数据集上实现的，不过实现的方式大同小异。无非就是图片输入的尺寸不一样。

---

一、model

这一部分实现每一层的的逻辑

import torch.nn as nn
import torch.nn.functional as F

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 5) # in_channels=3 out_channels=16 kernel=5
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
	# 调用上面定义的函数
    def forward(self, x):
        x = F.relu(self.conv1(x))    # input(3, 32, 32) output(16, 28, 28)
        x = self.pool1(x)            # output(16, 14, 14)
        x = F.relu(self.conv2(x))    # output(32, 10, 10)
        x = self.pool2(x)            # output(32, 5, 5)
        x = x.view(-1, 32*5*5)       # output(32*5*5)
        x = F.relu(self.fc1(x))      # output(120)
        x = F.relu(self.fc2(x))      # output(84)
        x = self.fc3(x)              # output(10)
        return x

二、train

这一部分开始对模型进行训练。

import torch
import torchvision
import torch.nn as nn

from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

def main():
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    # 50000张训练图片
    # 第一次使用时要将download设置为True才会自动去下载数据集
    train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
                                             download=False, transform=transform)
    train_loader = DataLoader(train_set, batch_size=36,
                                               shuffle=True, num_workers=0)

    # 10000张验证图片
    # 第一次使用时要将download设置为True才会自动去下载数据集
    val_set = torchvision.datasets.CIFAR10(root='./data', train=False,
                                           download=False, transform=transform)
    val_loader = DataLoader(val_set, batch_size=5000,
                                             shuffle=False, num_workers=0)
    val_data_iter = iter(val_loader)
    val_image, val_label = val_data_iter.next()

    # classes = ('plane', 'car', 'bird', 'cat',
    #            'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    # 定义模型和损失函数
    net = LeNet()
    loss_function = nn.CrossEntropyLoss()

    # 定义优化器
    optimizer = optim.Adam(net.parameters(),lr=0.001)

    # 对训练集迭代5次
    epochs = 5
    for epoch in range(epochs):
        running_loss = 0
        for step,data in enumerate(train_loader,start=0): # step从start开始
            inputs,labels = data
            # 梯度清零
            optimizer.zero_grad()

            outputs = net(inputs)
            # 计算损失函数
            loss = loss_function(outputs,labels)
            # 反向传播并更新参数
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if step%500 == 499:
                with torch.no_grad():
                    outputs = net(val_image)
                    predict_y = torch.max(outputs,dim=1)[1]
                    accuracy = (predict_y==val_label).sum().item() / val_label.size(0)
                    print('[%d  %3d] train_loss: %.3f test_accuracy: %.3f' %
                          (epoch+1,step+1,running_loss/500,accuracy))
                    running_loss = 0

    print('Finished Training')
    save_path = './Lenet.pth'
    torch.save(net.state_dict(),save_path)
if __name__ == '__main__':
    main()

三、test

import torch
import torchvision.transforms as transforms
from PIL import Image

from model import LeNet


def main():
    transform = transforms.Compose(
        [transforms.Resize((32, 32)), # 对输入的图片尺寸进行调整
         transforms.ToTensor(), # Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. (H x W x C)->(C x H x W)
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    
    net = LeNet()
    net.load_state_dict(torch.load('Lenet.pth')) # 加载模型

    im = Image.open('data/plane.jpg')
    im = transform(im)  # [C, H, W]
    # 增加一个维度，batch
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W] 增加一个batch维度

    with torch.no_grad():
        outputs = net(im)
        predict = torch.max(outputs, dim=1)[1].numpy()
    print(classes[int(predict)])

if __name__ == '__main__':
    main()

相关问题总结

1. with torch.no_grad():

with torch.no_grad():其实是一个上下文管理器，是的在with里的语句不再执行计算图的运算，所谓的计算图的运算，其实就是反向传播的时候需要。torch.no_grad() 会关闭自动求导引擎的， 因此能节省显存和加速。可以理解为被该语句 wrap 起来的部分将不会track 梯度。
但是加不加，计算结果都一样，只不过outputs少了一个属性。

2. torch.max()函数

torch.max(input, dim, keepdim=False, *, out=None)
可以点击链接查看pytorch的官方文档

输入：
input是softmax函数输出的一个tensor
dim是max函数索引的维度0/1，0是每列的最大值，1是每行的最大值

输出：
函数会返回两个tensor，第一个tensor是每行的最大值；第二个tensor是每行最大值的索引。

示例：就以这个模型的输出举例

训练时的输出示例：
我只在原来的代码上加入了输出语句。

with torch.no_grad():
    outputs = net(val_image)
    print(f'outputs的尺寸为：{outputs.size()}')
    print(f'outputs: {outputs}')
    print(f'torch.max(outputs,dim=1):{torch.max(outputs,dim=1)}')
    predict_y = torch.max(outputs,dim=1)[1]
    accuracy = (predict_y==val_label).sum().item() / val_label.size(0)

输出结果：

outputs其实可以理解为二维矩阵，但它严格意义上并不是二维矩阵，因为它是tensor类型的。
而torch.max(outputs,dim=1)其实是返回两个tensor，第一个tensor是每行的最大值；第二个tensor是每行最大值的索引。
这个索引就对应着classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')，这个顺序并不是瞎定的，是官方文档规定的。
那么重点来了，我们需要的是给你一个图片，需要我们看出这个图片到底是属于哪个类型，所以只有第二个tensor对我们有用。
所以我们这里使用predict_y = torch.max(outputs,dim=1)[1]，这也是[1]的来历。

在有的地方我们会看到torch.max(a, 1).data.numpy()的写法，这是因为在早期的pytorch的版本中，variable变量和tenosr是不一样的数据格式，variable可以进行反向传播，tensor不可以，需要将variable转变成tensor再转变成numpy。现在的版本已经将variable和tenosr合并，所以只用torch.max(a,1).numpy()就可以了。

那我们再来看一下在predict上的示例：

    with torch.no_grad():
        outputs = net(im)
        print(f'outputs: {outputs}')
        predict = torch.max(outputs, dim=1)[1].numpy()
        print(torch.max(outputs, dim=1))
        print(torch.max(outputs, dim=1)[0])
        print(torch.max(outputs, dim=1)[1])
    print(classes[int(predict)])

输出：

因为我们在测试的时候只是拿出一张图片，所以outputs的尺寸为（1,10），这个10就是有几种分类结果，对应class。
可以看到torch.max(outputs, dim=1)的输出，为两个tensor，我们只需要取出第二个就好。