PyTorch--残差网络（ResNet）在CIFAR-10数据集进行图像分类

完整代码

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms


# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Hyper-parameters
num_epochs = 80
batch_size = 100
learning_rate = 0.001

# Image preprocessing modules
transform = transforms.Compose([
    transforms.Pad(4),
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32),
    transforms.ToTensor()])

# CIFAR-10 dataset
train_dataset = torchvision.datasets.CIFAR10(root='../../data/',
                                             train=True, 
                                             transform=transform,
                                             download=True)

test_dataset = torchvision.datasets.CIFAR10(root='../../data/',
                                            train=False, 
                                            transform=transforms.ToTensor())

# Data loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)

# 3x3 convolution
def conv3x3(in_channels, out_channels, stride=1):
    return nn.Conv2d(in_channels, out_channels, kernel_size=3, 
                     stride=stride, padding=1, bias=False)

# Residual block
class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(ResidualBlock, self).__init__()
        self.conv1 = conv3x3(in_channels, out_channels, stride)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(out_channels, out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = downsample
        
    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.downsample:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out

# ResNet
class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=10):
        super(ResNet, self).__init__()
        self.in_channels = 16
        self.conv = conv3x3(3, 16)
        self.bn = nn.BatchNorm2d(16)
        self.relu = nn.ReLU(inplace=True)
        self.layer1 = self.make_layer(block, 16, layers[0])
        self.layer2 = self.make_layer(block, 32, layers[1], 2)
        self.layer3 = self.make_layer(block, 64, layers[2], 2)
        self.avg_pool = nn.AvgPool2d(8)
        self.fc = nn.Linear(64, num_classes)
        
    def make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if (stride != 1) or (self.in_channels != out_channels):
            downsample = nn.Sequential(
                conv3x3(self.in_channels, out_channels, stride=stride),
                nn.BatchNorm2d(out_channels))
        layers = []
        layers.append(block(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels
        for i in range(1, blocks):
            layers.append(block(out_channels, out_channels))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        out = self.relu(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.avg_pool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out
    
model = ResNet(ResidualBlock, [2, 2, 2]).to(device)


# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# For updating learning rate
def update_lr(optimizer, lr):    
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

# Train the model
total_step = len(train_loader)
curr_lr = learning_rate
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        
        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (i+1) % 100 == 0:
            print ("Epoch [{}/{}], Step [{}/{}] Loss: {:.4f}"
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

    # Decay learning rate
    if (epoch+1) % 20 == 0:
        curr_lr /= 3
        update_lr(optimizer, curr_lr)

# Test the model
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Accuracy of the model on the test images: {} %'.format(100 * correct / total))

# Save the model checkpoint
torch.save(model.state_dict(), 'resnet.ckpt')

这段代码是一个PyTorch实现的残差网络（ResNet），用于在CIFAR-10数据集上进行图像分类任务。下面是代码的详细解析：

代码解析

导入必要的库

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

导入PyTorch及其神经网络模块、torchvision库用于处理图像数据。

设备配置

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

设置运行设备，优先使用GPU，如果没有GPU，则使用CPU。

超参数设置

num_epochs = 80
batch_size = 100
learning_rate = 0.001

设置训练轮数、批次大小和学习率。

数据预处理

transform = transforms.Compose([
    transforms.Pad(4),
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32),
    transforms.ToTensor()])

定义数据预处理步骤，包括填充、随机水平翻转、随机裁剪和转换为张量。

加载CIFAR-10数据集

train_dataset = torchvision.datasets.CIFAR10(root='../../data/', train=True, transform=transform, download=True)
test_dataset = torchvision.datasets.CIFAR10(root='../../data/', train=False, transform=transforms.ToTensor())

加载CIFAR-10训练集和测试集。

创建数据加载器

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

创建用于加载数据的DataLoader。

定义3x3卷积函数

def conv3x3(in_channels, out_channels, stride=1):
    return nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)

定义一个3x3的卷积层。

定义残差块

class ResidualBlock(nn.Module):
    # ...

定义残差网络中的残差块，包含两个卷积层和批量归一化层。

定义ResNet模型

class ResNet(nn.Module):
    # ...

定义ResNet模型，使用残差块构建多个层。

实例化模型并移动到设备

model = ResNet(ResidualBlock, [2, 2, 2]).to(device)

创建ResNet模型实例并将其移动到配置的设备上。

定义损失函数和优化器

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

定义交叉熵损失函数和Adam优化器。

学习率衰减函数

def update_lr(optimizer, lr):    
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

定义一个函数用于更新优化器的学习率。

训练模型

for epoch in range(num_epochs):
    # ...
    # 每20个epoch衰减学习率
    if (epoch+1) % 20 == 0:
        curr_lr /= 3
        update_lr(optimizer, curr_lr)

执行训练循环，包括前向传播、损失计算、反向传播和参数更新，并每20个epoch衰减学习率。

测试模型

model.eval()
with torch.no_grad():
    # ...

在测试阶段，设置模型为评估模式，并计算准确率。

保存模型

torch.save(model.state_dict(), 'resnet.ckpt')

保存模型的状态字典。

这段代码实现了一个标准的ResNet架构，用于CIFAR-10数据集的分类任务。代码中包含了数据预处理、模型定义、训练过程、测试评估和模型保存等关键步骤。

常见函数及其用法

以下是代码中使用的常见函数及其解析：

1. torch.device

   • 格式：torch.device(device_str)
   • 参数：device_str —— 设备类型字符串（如’cuda’或’cpu’）。
   • 意义：确定模型和张量运行的设备。
   • 用法示例：device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

2. torchvision.datasets.CIFAR10

   • 格式：torchvision.datasets.CIFAR10(root, train, transform, download)
   • 参数：指定数据集的路径、是否为训练集、预处理变换、是否下载数据集。
   • 意义：加载CIFAR-10数据集。
   • 用法示例：train_dataset = torchvision.datasets.CIFAR10(root='../../data/', train=True, transform=transform, download=True)

3. torch.utils.data.DataLoader

   • 格式：torch.utils.data.DataLoader(dataset, batch_size, shuffle)
   • 参数：数据集对象、批次大小、是否打乱数据。
   • 意义：创建数据加载器，用于批量加载数据。
   • 用法示例：train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)

4. nn.Conv2d

   • 格式：nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
   • 参数：输入通道数、输出通道数、卷积核大小、步长、填充。
   • 意义：创建二维卷积层。
   • 用法示例：return nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)

5. nn.BatchNorm2d

   • 格式：nn.BatchNorm2d(num_features)
   • 参数：特征数量。
   • 意义：创建二维批量归一化层。
   • 用法示例：self.bn1 = nn.BatchNorm2d(out_channels)

6. nn.ReLU

   • 格式：nn.ReLU(inplace=True/False)
   • 参数：是否使用内存原地(inplace)优化。
   • 意义：创建ReLU激活层。
   • 用法示例：self.relu = nn.ReLU(inplace=True)

7. nn.Sequential

   • 格式：nn.Sequential(*modules)
   • 参数：一个模块序列。
   • 意义：按顺序应用多个模块。
   • 用法示例：downsample = nn.Sequential(conv3x3(self.in_channels, out_channels, stride=stride), nn.BatchNorm2d(out_channels))

8. nn.CrossEntropyLoss

   • 格式：nn.CrossEntropyLoss()
   • 意义：创建交叉熵损失层，用于多分类问题。
   • 用法示例：criterion = nn.CrossEntropyLoss()

9. torch.optim.Adam

   • 格式：torch.optim.Adam(params, lr)
   • 参数：模型参数、学习率。
   • 意义：创建Adam优化器。
   • 用法示例：optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

10. .to(device)

    • 格式：.to(device)
    • 参数：设备对象。
    • 意义：将模型或张量移动到指定设备。
    • 用法示例：images = images.to(device)

11. view

    • 格式：view(size)
    • 参数：新的大小。
    • 意义：重塑张量。
    • 用法示例：out = out.view(out.size(0), -1)

12. max

    • 格式：max(dim, keepdim)
    • 参数：计算最大值的维度、是否保持维度。
    • 意义：计算并返回张量在指定维度上的最大值和索引。
    • 用法示例：_, predicted = torch.max(outputs.data, 1)

13. no_grad

    • 格式：torch.no_grad()
    • 意义：上下文管理器，用于禁用梯度计算。
    • 用法示例：with torch.no_grad():

14. sum

    • 格式：sum(dim, keepdim)
    • 参数：求和的维度、是否保持维度。
    • 意义：计算张量在指定维度的和。
    • 用法示例：correct += (predicted == labels).sum().item()

15. torch.save

    • 格式：torch.save(obj, f)
    • 参数：要保存的对象、文件路径。
    • 意义：保存对象到文件。
    • 用法示例：torch.save(model.state_dict(), 'resnet.ckpt')

这些函数和类是构建、训练和测试PyTorch模型的基础，涵盖了设备配置、数据加载、模型定义、训练过程、测试评估和模型保存等关键步骤。

运行过程

在整体进行一些可视化改进之后，可以看到效果图如下图所示：