LeNet Pytorch实现【初学者】

1. 内容来源

23 经典卷积神经网络 LeNet【动手学深度学习v2】_哔哩哔哩_bilibili23 经典卷积神经网络 LeNet【动手学深度学习v2】共计3条视频，包括：LeNet、代码、QA等，UP主更多精彩视频，请关注UP账号。https://www.bilibili.com/video/BV1t44y1r7ct/?spm_id_from=333.999.0.0&vd_source=f94822d3eca79b8e245bb58bbced6b77

2. Pytorch实现

2.1 模型构建

由于此篇文章（LeCun Y, Bottou L, Bengio Y, et al. Gradient-based learning applied to document recognition[J]. Proceedings of the IEEE, 1998, 86(11): 2278-2324.）是历史上第一个深度学习网络，原文中某些操作与现在常用的层有少许出入，沐神简化成以下模型：

1. 将数据集reshape成(batch_size, 1， 28， 28）

2. 经过卷积层（输出6通道，使用5*5卷积核，2填充）-- 激活层（Sigmoid）-- 平均池化层（使用2*2核，2步长）

3. 经过卷积层（输出16通道，使用5*5卷积核，无填充）-- 激活层（Sigmoid）-- 平均池化层（使用2*2核，2步长）

4. 展开成一维向量特征，经过线性层（输入维度16*5*5，输出维度120）-- 激活层（Sigmoid）

5. 经过线性层（输入维度120，输出维度84）-- 激活层（Sigmoid）

6. 输出层（输入维度84，输出维度10）

代码实现：

import torch
from torch import nn
from d2l import torch as d2l

class Reshape(torch.nn.Module):
    def forward(self, x):
        return x.view(-1, 1, 28, 28)
    
net = torch.nn.Sequential(
    Reshape(), 
    nn.Conv2d(1, 6, kernel_size=5, padding=2), 
    nn.Sigmoid(), 
    nn.AvgPool2d(kernel_size=2, stride=2), 
    nn.Conv2d(6, 16, kernel_size=5), 
    nn.Sigmoid(), 
    nn.AvgPool2d(kernel_size=2, stride=2), 
    nn.Flatten(), 
    nn.Linear(16 * 5 * 5, 120), 
    nn.Sigmoid(), 
    nn.Linear(120, 84), 
    nn.Sigmoid(), 
    nn.Linear(84, 10)
)

其中第一个线性层的16*5*5可以手推，也可以使用以下方式打印每一层的输出维度

代码实现：

# Dimension Check
X = torch.rand(size=(1, 1, 28, 28), dtype=torch.float32)
for layer in net:
    X = layer(X)
    print(layer.__class__.__name__, 'output shape:    \t', X.shape)

输出：

Reshape output shape:    	 torch.Size([1, 1, 28, 28])
Conv2d output shape:    	 torch.Size([1, 6, 28, 28])
Sigmoid output shape:    	 torch.Size([1, 6, 28, 28])
AvgPool2d output shape:    	 torch.Size([1, 6, 14, 14])
Conv2d output shape:    	 torch.Size([1, 16, 10, 10])
Sigmoid output shape:    	 torch.Size([1, 16, 10, 10])
AvgPool2d output shape:    	 torch.Size([1, 16, 5, 5])
Flatten output shape:    	 torch.Size([1, 400])
Linear output shape:    	 torch.Size([1, 120])
Sigmoid output shape:    	 torch.Size([1, 120])
Linear output shape:    	 torch.Size([1, 84])
Sigmoid output shape:    	 torch.Size([1, 84])
Linear output shape:    	 torch.Size([1, 10])

2.2 模型训练

由于从LeNet开始网络训练计算量开始指数型暴涨，需要借助GPU进行加速，所以重写train函数

代码实现：

def train_gpu(net, train_iter, test_iter, num_epochs, lr, device):
    # Parameter Initialization
    def init_weight(m):
        if type(m) == nn.Linear or type(m) == nn.Conv2d:
            nn.init.xavier_uniform_(m.weight)
    net.apply(init_weight)
    # Device Choosing
    print('training on', device)
    net.to(device)
    # Trainging
    optimizer = torch.optim.SGD(net.parameters(), lr=lr)
    loss = nn.CrossEntropyLoss()
    animator = d2l.Animator(xlabel='Epoch', 
                            xlim=[1, num_epochs], 
                            legend=['Train Loss', 'Train Acc', 'Test Acc']
                            )
    timer, num_batches = d2l.Timer(), len(train_iter)
    for epoch in range(num_epochs):
        metric = d2l.Accumulator(3)
        net.train()
        for i, (X, y) in enumerate(train_iter):
            timer.start()
            optimizer.zero_grad()
            X, y = X.to(device), y.to(device)
            yhat = net(X)
            l = loss(yhat, y)
            l.backward()
            optimizer.step()
            metric.add(l * X.shape[0], d2l.accuracy(yhat, y), X.shape[0])
            timer.stop()
            train_l = metric[0] / metric[2]
            train_acc = metric[1] / metric[2]
            if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:
                animator.add(epoch + (1 + i) / num_batches, 
                             (train_l, train_acc, None))
        test_acc = evaluate_accuracy_gpu(net, test_iter)
        animator.add(epoch + 1, (None, None, test_acc))
    print(f'loss {train_l:.3f}, train acc {train_acc:.3f}, test acc {test_acc:.3f}')
    print(f'{metric[2] * num_epochs / timer.sum():.1f} examples/sec', 
          f'on {str(device)}')

其中，涉及matric、animator代码是沐神用来动态展示训练过程的可以省略，与在CPU上训练不同的在于所有涉及计算的数据均需要使用.to(device)搬运到GPU上，与模型所在位置保持一致。device可以用沐神封装的try_gpu()函数自动获取名称。

try_gpu()源码：

def try_gpu(i=0):
    """Return gpu(i) if exists, otherwise return cpu().

    Defined in :numref:`sec_use_gpu`"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

 训练代码：

#Training
lr, num_epochs = 0.9, 10
train_gpu(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())

结果：