PyTorch学习笔记——ResNet的实现

Posted 2022-05-18 raelum

目录

• 一、Residual Block
• 二、ResNet 架构
• 三、训练/测试 ResNet
• 附录一：完整代码
• 附录二：ResNet 在 FashionMNIST 数据集上的表现

一、Residual Block

残差块有以下两种：

实现如下：

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


class Residual(nn.Module):

    def __init__(self, in_channels, out_channels, stride=1, conv_1x1=False):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, stride=stride),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
        )
        self.conv_1x1 = nn.Conv2d(in_channels, out_channels, kernel_size=1,
                                  stride=stride) if conv_1x1 else None

    def forward(self, x):
        y = self.block(x)
        if self.conv_1x1:
            x = self.conv_1x1(x)
        return F.relu(y + x)

二、ResNet 架构

从上图可以看出，该架构有 17 个卷积层和 1 个全连接层。因此，这种模型通常被称为 ResNet-18。虽然ResNet的主体架构跟GoogLeNet类似，但ResNet架构更简单，修改也更方便。这些因素都导致了ResNet迅速被广泛使用。

ResNet-18 的实现如下：

class ResNet(nn.Module):

    def __init__(self):
        super().__init__()
        self.block_1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
        )
        self.block_2 = nn.Sequential(
            Residual(64, 64),
            Residual(64, 64),
            Residual(64, 128, stride=2, conv_1x1=True),
            Residual(128, 128),
            Residual(128, 256, stride=2, conv_1x1=True),
            Residual(256, 256),
            Residual(256, 512, stride=2, conv_1x1=True),
            Residual(512, 512),
        )
        self.block_3 = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.block_1(x)
        x = self.block_2(x)
        x = self.block_3(x)
        return x

三、训练/测试 ResNet

我们使用 CIFAR-10 数据集，设置学习率大小为 0.05，batch size 为 128，在 NVIDIA GeForce RTX 3080 Ti 训练 20 个 Epoch 的结果如下：

Epoch 20
--------------------------------------------------
Train Avg Loss: 0.000445, Train Accuracy: 1.000000
Test  Avg Loss: 0.759453, Test  Accuracy: 0.824000

--------------------------------------------------
3163.0 samples/sec
--------------------------------------------------

Done!

附录一：完整代码

import torchvision
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor, Resize
from Experiment import Experiment as E


class Residual(nn.Module):

    def __init__(self, in_channels, out_channels, stride=1, conv_1x1=False):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, stride=stride),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
        )
        self.conv_1x1 = nn.Conv2d(in_channels, out_channels, kernel_size=1,
                                  stride=stride) if conv_1x1 else None

    def forward(self, x):
        y = self.block(x)
        if self.conv_1x1:
            x = self.conv_1x1(x)
        return F.relu(y + x)


class ResNet(nn.Module):

    def __init__(self):
        super().__init__()
        self.block_1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
        )
        self.block_2 = nn.Sequential(
            Residual(64, 64),
            Residual(64, 64),
            Residual(64, 128, stride=2, conv_1x1=True),
            Residual(128, 128),
            Residual(128, 256, stride=2, conv_1x1=True),
            Residual(256, 256),
            Residual(256, 512, stride=2, conv_1x1=True),
            Residual(512, 512),
        )
        self.block_3 = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.block_1(x)
        x = self.block_2(x)
        x = self.block_3(x)
        return x


def init_net(m):
    if type(m) == nn.Linear or type(m) == nn.Conv2d:
        nn.init.xavier_uniform_(m.weight)


transformer = torchvision.transforms.Compose([Resize(224), ToTensor()])
train_data = torchvision.datasets.CIFAR10('/mnt/mydataset',
                                          train=True,
                                          transform=transformer,
                                          download=True)
test_data = torchvision.datasets.CIFAR10('/mnt/mydataset',
                                         train=False,
                                         transform=transformer,
                                         download=True)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True, num_workers=4)
test_loader = DataLoader(test_data, batch_size=128, num_workers=4)

resnet = ResNet()
resnet.apply(init_net)
e = E(train_loader, test_loader, resnet, 20, 0.05)
e.main()
e.show()

附录二：ResNet 在 FashionMNIST 数据集上的表现

设置 Batch size 为 128，学习率为 0.03，只训练 10 个 Epoch，ResNet 在 FashionMNIST 数据集上的表现如下所示（GPU 为 NVIDIA RTX A6000）：

Epoch 10
--------------------------------------------------
Train Avg Loss: 0.018614, Train Accuracy: 0.995433
Test  Avg Loss: 0.285646, Test  Accuracy: 0.920600

--------------------------------------------------
5093.3 samples/sec
--------------------------------------------------

Done!