pytorch实现简单的ResNet并对MNIST进行分类

问题说明

• 利用pytorch实现简单的ResNet（2个ResidualBlock）
• 以MNIST数据集为例进行分类任务

模块设计

1.ResidualBlock的结构

本例中实现的ResidualBlock的结构如下图所示：

其中Weight Layer是一个卷积层。这个ResidualBlock的一个特点是在最后进行第二次的激活时，将第二次卷积的结果加上原始的ResidualBlock的输入x，最后再进行relu激活。
原因：主要是为了避免梯度消失。在进行求导时，d(H(x))=d(F(x)) + 1，保证最后的导数不会趋近于0.

• 实现代码如下：

class ResidualBlock(nn.Module):
    """
    每一个ResidualBlock,需要保证输入和输出的维度不变
    所以卷积核的通道数都设置成一样
    """
    def __init__(self, channel):
        super().__init__()
        self.conv1 = nn.Conv2d(channel, channel, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(channel, channel, kernel_size=3, padding=1)

    def forward(self, x):
        """
        ResidualBlock中有跳跃连接;
        在得到第二次卷积结果时,需要加上该残差块的输入,
        再将结果进行激活,实现跳跃连接 ==> 可以避免梯度消失
        在求导时,因为有加上原始的输入x,所以梯度为: dy + 1,在1附近
        """
        y = F.relu(self.conv1(x))
        y = self.conv2(y)

        return F.relu(x + y)

2.ResNet的简单结构

本例中实现的ResNet结构如下图所示：

说明：如上图所示，本例中的网络结构包括：2个卷积层、2个池化层、2个激活层、2个ResidualBlock，1个全连接层。

代码实现

• 完整代码如下：

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
 @Author   : Sword
 @Date     : 2020/5/22 0022 
 @Time     : 下午 15:06
 @Version  : v1.0
 @File     : ResNet_mnist.py 
 @Describe :ResNet分类MNIST
"""
# 这个是python当中让print都以python3的形式进行print,即把print视为函数
from __future__ import print_function
# 使得我们能够手动输入命令行参数
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable


# Training settings
# 设置一些参数,每个都有默认值,输入python main.py -h可以获得帮助
parser = argparse.ArgumentParser(description='Pytorch MNIST Example')
parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                    help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                    help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=10, metavar='N',
                    help='number of epochs to train (default: 10')
parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                    help='learning rate (default: 0.01)')
parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                    help='SGD momentum (default: 0.5)')
parser.add_argument('--no-cuda', action='store_true', default=True,
                    help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
# 跑多少次batch进行一次日志记录
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                    help='how many batches to wait before logging training status')
# 这个是使用argparse模块时的必备行,将参数进行关联
args = parser.parse_args()
# 这个是在确认是否使用GPU的参数
args.cuda = not args.no_cuda and torch.cuda.is_available()

# 设置一个随机数种子
torch.manual_seed(args.seed)
if args.cuda:
    # 为GPU设置一个随机数种子
    torch.cuda.manual_seed(args.seed)

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307, ), (0.3081, ))
])

train_set = datasets.MNIST(root='../dataset/mnist', train=True, transform=transform, download=False)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True)
test_set = datasets.MNIST(root='../dataset/mnist', train=False, transform=transform, download=False)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=args.test_batch_size, shuffle=False)


class ResidualBlock(nn.Module):
    """
    每一个ResidualBlock,需要保证输入和输出的维度不变
    所以卷积核的通道数都设置成一样
    """
    def __init__(self, channel):
        super().__init__()
        self.conv1 = nn.Conv2d(channel, channel, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(channel, channel, kernel_size=3, padding=1)

    def forward(self, x):
        """
        ResidualBlock中有跳跃连接;
        在得到第二次卷积结果时,需要加上该残差块的输入,
        再将结果进行激活,实现跳跃连接 ==> 可以避免梯度消失
        在求导时,因为有加上原始的输入x,所以梯度为: dy + 1,在1附近
        """
        y = F.relu(self.conv1(x))
        y = self.conv2(y)

        return F.relu(x + y)


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=5)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5)
        self.res_block_1 = ResidualBlock(16)
        self.res_block_2 = ResidualBlock(32)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(512, 10)

    def forward(self, x):
        in_size = x.size(0)
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = self.res_block_1(x)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = self.res_block_2(x)
        x = x.view(in_size, -1)
        x = self.fc1(x)
        return F.log_softmax(x, dim=1)


model = Net()

# 判断是否调用GPU模式
if args.cuda:
    model.cuda()
# 初始化优化器 model.train()
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)


def train(epoch):
    """
    定义每个epoch的训练细节
    """
    # 设置为training模式
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # 如果要调用GPU模式,就把数据转存到GPU
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)

        # 优化器梯度初始化为零
        optimizer.zero_grad()
        output = model(data)
        # 负对数似然函数损失
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tloss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()
            ))


def test():
    # 设置为test模式
    model.eval()
    # 初始化测试损失值为0
    test_loss = 0
    # 初始化预测正确的数据个数为0
    correct = 0
    for data, target in test_loader:
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)
        output = model(data)
        # 把所有loss值进行累加
        test_loss += F.nll_loss(output, target, size_average=False).item()
        # 获取最大对数概率值的索引
        pred = output.data.max(1, keepdim=True)[1]
        # 对预测正确的个数进行累加
        correct += pred.eq(target.data.view_as(pred)).cpu().sum()

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


# 进行每个epoch的训练
for epoch in range(1, args.epochs + 1):
    train(epoch)
    test()

• 实验结果
  本例中只跑了10个epoch，最终在测试集上能达到99%的精确度。
  
• 说明：上述图片截图自哔哩哔哩的UP主（@刘二大人）的视频。