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 = downsampledef forward(self, x):residual = xout = 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 += residualout = 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 = 16self.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 = Noneif (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_channelsfor 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 outmodel = 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 passoutputs = model(images)loss = criterion(outputs, labels)# Backward and optimizeoptimizer.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 rateif (epoch+1) % 20 == 0:curr_lr /= 3update_lr(optimizer, curr_lr)# Test the model
model.eval()
with torch.no_grad():correct = 0total = 0for 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 /= 3update_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模型的基础，涵盖了设备配置、数据加载、模型定义、训练过程、测试评估和模型保存等关键步骤。

运行过程

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