Pytorch学习笔记：CIFAR10彩色图片识别

• 🍨 本文为🔗365天深度学习训练营 中的学习记录博客
• 🍖 原作者：K同学啊

一、实验目的

了解PyTorch的基本用法，并用它搭建一个深度学习程序，用来完成CIFAR10彩色图片识别。

二、课程分析

P2周课程与P1对比：
1. “练习册”不一样了：
• P1 (MNIST)： 是 28x28 的黑白数字图片 (1 个通道)。
• P2 (CIFAR10)： 是 32x32 的彩色物体图片 (比如飞机、小鸟、卡车)。因为是彩色的，所以有 3 个通道（红、绿、蓝）。
2. “大脑”变得更“大”了：
• 因为“彩色物体”比“黑白数字”复杂得多，所以我们 P2 的“大脑”(class Model)需要更“聪明”、更“深”。
• P1 的“大脑”只有一层“眼睛”(conv1) 和一层“决策官”(fc1)。
• P2 的“大脑”有两层“眼睛”(conv1, conv2) 和三层“决策官”(fc1, fc2, fc3)。
3. “流程”完全一样：
• 除了“大脑”的“设计图” 不一样，剩下的所有流程逻辑和P1是一样的：
• DataLoader 打包数据
• def train(…) 学习流程
• def test(…) 考试流程
• plt.plot(…) 画结果曲线

三、核心代码

1.前期准备

# (和 P1 基本一样)
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# 设置GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

2.导入数据 (下载并打包 CIFAR10)

# P2 的数据预处理(transform)更复杂，因为是彩色图
# 我们需要 "Normalize" (标准化)，让电脑看得更舒服
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# batch_size 我们用 32 (课程里好像是4，但 32 更快)
batch_size = 32

# 下载 P2 的“练习册” torchvision.datasets.CIFAR10
train_ds = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
train_dl = torch.utils.data.DataLoader(train_ds, batch_size=batch_size,
                                          shuffle=True)

# 下载 P2 的“考试卷”
test_ds = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
test_dl = torch.utils.data.DataLoader(test_ds, batch_size=batch_size,
                                         shuffle=False)a

# P2 的 10 个分类名字
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

3.数据可视化

print("开始画图，请稍等...")

# 先从"练习册"里拿出一份(batch)数据
imgs_batch, labels_batch = next(iter(train_dl))

# (画图的代码，P2 和 P1 有点不一样，因为是彩色图)
def imshow(img):
    img = img / 2 + 0.5     # "反标准化" (unnormalize)
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.axis('off') # 不显示坐标轴

# 指定图片大小
plt.figure(figsize=(20, 5))
# 画出前 20 张图
imshow(torchvision.utils.make_grid(imgs_batch[:20]))
# 打印这 20 张图的“标准答案”
print(' '.join(f'{classes[labels_batch[j]]:5s}' for j in range(20)))
# 秀出图片
plt.show()

print("画图结束，继续训练...")

4.构建神经网络

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # “眼睛” (第1层)
        # (注意：P1 是 1 通道输入，P2 是 3 通道输入)
        self.conv1 = nn.Conv2d(3, 16, kernel_size=5) # 3(彩色) -> 16
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # “眼睛” (第2层)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5) # 16 -> 32
        
        self.relu = nn.ReLU() # 激活函数 (P1 P2 一样)
        
        # “决策官” (有3层)
        self.fc1 = nn.Linear(32 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10) # 最后还是输出 10 个“分数”

    def forward(self, x):
        # 数据在“大脑”里的走法
        x = self.conv1(x) # 走“眼睛1”
        x = self.relu(x)
        x = self.pool(x)
        
        x = self.conv2(x) # 走“眼睛2”
        x = self.relu(x)
        x = self.pool(x)
        
        x = x.view(-1, 32 * 5 * 5) # “摊平”
        
        x = self.fc1(x) # 走“决策官1”
        x = self.relu(x)
        x = self.fc2(x) # 走“决策官2”
        x = self.relu(x)
        x = self.fc3(x) # 走“决策官3”，得到10个“分数”
        return x

# 把我们搭好的“大脑”搬到GPU上
model = Model().to(device)

5. 定义“评分标准”和“优化方法”

# (和 P1 一样)
learning_rate = 0.001
epochs = 5 # (P2 的大脑更“大”，可以先设成 5 跑跑看)

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

6. 编写训练和测试函数

# (和 P1 一样)
#  定义一个“学习”的流程
def train(dataloader, model, loss_fn, optimizer):
    model.train() 
    for x, y in dataloader:
        x, y = x.to(device), y.to(device)
        pred = model(x)
        loss = loss_fn(pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# 定义一个“考试”的流程
def test(dataloader, model, loss_fn):
    model.eval() 
    test_loss, correct = 0, 0
    with torch.no_grad():
        for x, y in dataloader:
            x, y = x.to(device), y.to(device)
            pred = model(x) 
            test_loss += loss_fn(pred, y).item() 
            correct += (pred.argmax(1) == y).type(torch.float).sum().item() 
    test_loss /= len(dataloader) 
    correct /= len(dataloader.dataset) 
    return test_loss, correct  # 返回“罚分”和“正确率”

7. 正式训练（循环执行）

# (和 P1 一样)
print("开始正式训练...")

train_loss_history = []
train_acc_history = []
test_loss_history = []
test_acc_history = []

for epoch in range(epochs):
    print(f"Epoch {epoch+1}\n-------------------------------")
    train(train_dl, model, loss_fn, optimizer)
    
    train_loss, train_acc = test(train_dl, model, loss_fn)
    test_loss, test_acc = test(test_dl, model, loss_fn)
    
    train_loss_history.append(train_loss)
    train_acc_history.append(train_acc)
    test_loss_history.append(test_loss)
    test_acc_history.append(test_acc)
    
    print(f"Train Error: \n Accuracy: {(100*train_acc):>0.1f}%, Avg loss: {train_loss:>8f}")
    print(f"Test Error: \n Accuracy: {(100*test_acc):>0.1f}%, Avg loss: {test_loss:>8f} \n")

print("Done!")

8.结果可视化

# (和 P1 一样)
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(range(epochs), train_loss_history, label='Training Loss')
plt.plot(range(epochs), test_loss_history, label='Test Loss')
plt.title('Training and Test Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend() 

plt.subplot(1, 2, 2)
plt.plot(range(epochs), train_acc_history, label='Training Accuracy')
plt.plot(range(epochs), test_acc_history, label='Test Accuracy')
plt.title('Training and Test Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend() 

plt.show()

结果：

四、总结