实训p66_2

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import Compose, ToTensor, Normalize
import matplotlib.pyplot as plt
import numpy as np

定义数据预处理流程：转换为张量 + 归一化（CIFAR-10均值和标准差统计值）

transform = Compose([
ToTensor(),
Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010])
])

加载训练集和测试集

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

构建数据加载器（批量加载数据，加速训练）

batch_size = 128
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

类别标签映射

classes = ('飞机', '汽车', '鸟', '猫', '鹿', '狗', '青蛙', '马', '船', '卡车')
class CIFAR10Model(nn.Module):
def init(self):
super(CIFAR10Model, self).init()
# 第一层卷积：3通道输入，32个卷积核，卷积核大小3x3，激活函数ReLU
self.conv1 = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2) # 池化层：缩小特征图尺寸
)
# 第二层卷积
self.conv2 = nn.Sequential(
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
# 第三层卷积
self.conv3 = nn.Sequential(
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
# 全连接层：将卷积特征展平后映射到类别
self.fc = nn.Sequential(
nn.Flatten(), # 展平特征图
nn.Linear(128 * 4 * 4, 512), # 全连接层1
nn.ReLU(),
nn.Dropout(0.5), # Dropout防止过拟合
nn.Linear(512, 10) # 全连接层2：输出10个类别
)

def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)
    x = self.conv3(x)
    x = self.fc(x)
    return x

初始化模型、损失函数和优化器

model = CIFAR10Model()
criterion = nn.CrossEntropyLoss() # 分类任务常用交叉熵损失
optimizer = optim.Adam(model.parameters(), lr=0.001) # Adam优化器
epochs = 50
train_losses = []
train_accs = []
test_losses = []
test_accs = []

for epoch in range(epochs):
# 训练阶段
model.train()
train_loss = 0.0
correct = 0
total = 0

for inputs, targets in train_loader:
    optimizer.zero_grad()  # 清空梯度
    outputs = model(inputs)  # 前向传播
    loss = criterion(outputs, targets)  # 计算损失
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数
    
    train_loss += loss.item()
    _, predicted = outputs.max(1)
    total += targets.size(0)
    correct += predicted.eq(targets).sum().item()

train_loss /= len(train_loader)
train_acc = 100. * correct / total
train_losses.append(train_loss)
train_accs.append(train_acc)

# 测试阶段
model.eval()
test_loss = 0.0
correct = 0
total = 0

with torch.no_grad():  # 测试时无需计算梯度
    for inputs, targets in test_loader:
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
        test_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

test_loss /= len(test_loader)
test_acc = 100. * correct / total
test_losses.append(test_loss)
test_accs.append(test_acc)

print(f'Epoch {epoch+1}/{epochs} | '
      f'Train Loss: {train_loss:.3f} | Train Acc: {train_acc:.2f}% | '
      f'Test Loss: {test_loss:.3f} | Test Acc: {test_acc:.2f}%')

保存模型

torch.save(model.state_dict(), 'cifar10_model.pth')
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Train Loss')
plt.plot(test_losses, label='Test Loss')
plt.legend()
plt.title('Loss Curve')

plt.subplot(1, 2, 2)
plt.plot(train_accs, label='Train Accuracy')
plt.plot(test_accs, label='Test Accuracy')
plt.legend()
plt.title('Accuracy Curve')
plt.show()

测试集整体精度

model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, targets in test_loader:
outputs = model(inputs)
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
print(f'最终测试集精度: {100. * correct / total:.2f}%')

单张图片预测示例

def predict_single_image(image_index):
model.eval()
img, label = test_dataset[image_index]
with torch.no_grad():
output = model(img.unsqueeze(0)) # 增加批次维度
_, predicted = output.max(1)
print(f'真实类别: {classes[label]}')
print(f'预测类别: {classes[predicted.item()]}')
# 可视化图片
img = img.permute(1, 2, 0) # 转换维度：(C,H,W) → (H,W,C)
img = img * torch.tensor([0.2023, 0.1994, 0.2010]) + torch.tensor([0.4914, 0.4822, 0.4465]) # 反归一化
img = np.clip(img.numpy(), 0, 1)
plt.imshow(img)
plt.show()

预测第5张测试图

predict_single_image(5)