本项目实现了两种不同的神经网络模型在CIFAR-10数据集上的图像分类任务，包括线性分类器和全连接神经网络(FCNN)，并比较了不同优化器和学习率调度器的性能。😋
基于PyTorch的CIFAR-10图像分类实践
不得不说这次作业是最简单的一次，只需要简单地调一下包就可以轻松实现。
另一方面也展现出神经网络在 CV 方面的强大实力。😇
核心代码解析
1. 数据预处理与加载
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
# 数据预处理：归一化到[-1,1]范围
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# CIFAR-10类别标签
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
batch_size = 10
2. 模型定义
线性分类器 (LinearClassifier)
class LinearClassifier(nn.Module):
def __init__(self, in_channels: int, out_channels: int):
super().__init__()
# 将3x32x32的图像展平为3072维向量
self.linear = nn.Linear(in_features=in_channels, out_features=out_channels)
def forward(self, x: torch.Tensor):
x = x.view(x.size(0), -1) # 展平输入
return self.linear(x)
全连接神经网络 (FCNN)
class FCNN(nn.Module):
def __init__(self, in_channels: int, hidden_channels: int, out_channels: int):
super().__init__()
# 使用卷积层提取特征
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
# 全连接层进行分类
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
self.dropout = nn.Dropout(0.5)
def forward(self, x: torch.Tensor):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
3. 训练流程
def train(model, optimizer, scheduler, args):
# 加载训练数据
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
shuffle=True, num_workers=2)
EPOCHS = 5
train_losses = []
train_accuracies = []
for epoch in range(EPOCHS):
running_loss = 0.0
for i, data in enumerate(tqdm(trainloader)):
inputs, labels = data
# 前向传播
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播
loss.backward()
optimizer.step()
# 记录训练指标
running_loss += loss.item()
# ... 准确率计算代码
4. 测试与评估
def test(model, args):
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
shuffle=False, num_workers=2)
# 加载训练好的模型
model.load_state_dict(torch.load(PATH))
# 在测试集上评估
evaluate(model, testloader)
def evaluate(model, testloader):
model.eval()
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f'Test Accuracy: {accuracy:.2f}%')
return accuracy
5. 配置与命令行接口
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='The configs')
parser.add_argument('--run', type=str, default='train')
parser.add_argument('--model', type=str, default='linear')
parser.add_argument('--optimizer', type=str, default='adamw')
parser.add_argument('--scheduler', type=str, default='step')
args = parser.parse_args()
# 根据参数选择模型、优化器和调度器
# ... 配置代码
使用方法
训练模型
# 使用线性分类器 + AdamW优化器 + StepLR调度器
python main.py --run=train --model=linear --optimizer=adamw --scheduler=step
# 使用FCNN + SGD优化器 + Cosine调度器
python main.py --run=train --model=fcnn --optimizer=sgd --scheduler=cosine
测试模型
python main.py --run=test --model=fcnn
实验比较
优化器比较 (基于FCNN)
- AdamW: 自适应学习率，训练稳定，收敛快
- SGD: 需要手动调参，可能找到更优解
学习率调度器比较 (基于FCNN)
- StepLR: 在指定epoch降低学习率
- CosineAnnealingLR: 余弦退火，平滑调整学习率
可视化结果
项目包含完整的训练曲线可视化功能，可以显示：
- 每个batch的loss和accuracy变化
- 每个epoch的平均loss和accuracy
- 测试集上的分类结果可视化
技术要点
- 数据预处理: 图像归一化处理
- 模型设计: 从简单线性模型到复杂卷积网络
- 训练技巧: 梯度下降、学习率调度、dropout正则化
- 评估方法: 准确率计算、损失函数监控
- 可视化: 训练过程监控和结果展示