深度学习篇---PyTorch 图像分类模型从零 jupyter notebook 实现

Ronin-Lotus

953人浏览 · 2025-09-07 23:25:46

Ronin-Lotus · 2025-09-07 23:25:46 发布

总体流程：

下面我将为您提供一个完整的 PyTorch 图像分类教程，适合在 Jupyter Notebook 中逐步实现。我们将使用 CIFAR-10 数据集，从零开始构建一个卷积神经网络 (CNN)，并详细解释每一步的操作和输出含义。

步骤 1：准备工作与环境配置

首先，我们需要导入必要的库并设置随机种子以确保结果可复现。

# 导入必要的库
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np

# 设置随机种子，确保结果可复现
torch.manual_seed(42)
np.random.seed(42)

# 检查是否有可用的GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用的设备: {device}")

输出解释：

如果您的电脑有 NVIDIA 显卡并正确安装了 CUDA，输出会显示 "使用的设备: cuda"
否则会显示 "使用的设备: cpu"
使用 GPU 可以显著加速模型训练过程

步骤 2：数据加载与预处理

接下来，我们加载 CIFAR-10 数据集并进行必要的预处理。

# 定义数据变换
transform = transforms.Compose([
    transforms.Resize((32, 32)),  # 确保图像大小一致
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化
])

# 加载训练集和测试集
train_dataset = datasets.CIFAR10(
    root='./data', 
    train=True, 
    download=True, 
    transform=transform
)

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

# 创建数据加载器
batch_size = 64
train_loader = DataLoader(
    train_dataset, 
    batch_size=batch_size, 
    shuffle=True, 
    num_workers=2
)

test_loader = DataLoader(
    test_dataset, 
    batch_size=batch_size, 
    shuffle=False, 
    num_workers=2
)

# CIFAR-10的类别名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

# 查看数据集基本信息
print(f"训练集大小: {len(train_dataset)}")
print(f"测试集大小: {len(test_dataset)}")
print(f"批次大小: {batch_size}")
print(f"训练批次数量: {len(train_loader)}")
print(f"测试批次数量: {len(test_loader)}")

输出解释：

首次运行时会下载 CIFAR-10 数据集，大约 170MB
输出应显示：训练集大小: 50000，测试集大小: 10000
训练批次数量: 782 (50000/64 ≈ 782)
测试批次数量: 157 (10000/64 ≈ 157)

步骤 3：数据可视化

让我们可视化一些训练数据，了解我们要处理的图像是什么样的。

# 函数：显示图像
def imshow(img):
    img = img / 2 + 0.5  # 反标准化
    npimg = img.numpy()
    plt.figure(figsize=(10, 4))
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.axis('off')
    plt.show()

# 获取一些随机的训练图像
dataiter = iter(train_loader)
images, labels = next(dataiter)

# 显示图像
imshow(torchvision.utils.make_grid(images[:5]))
# 打印标签
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(5)))

输出解释：

会显示 5 张随机的 CIFAR-10 图像
图像下方会显示对应的类别名称
这些图像是 32x32 像素的彩色图像 (RGB)
您可能会看到类似 "cat ship ship plane frog" 这样的标签输出

步骤 4：定义卷积神经网络模型

现在我们来定义一个卷积神经网络模型。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 第一个卷积块
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.pool = nn.MaxPool2d(2, 2)
        
        # 第二个卷积块
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        
        # 第三个卷积块
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        
        # 全连接层
        self.fc1 = nn.Linear(128 * 4 * 4, 512)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        # 第一个卷积块: 卷积 -> 批归一化 -> ReLU -> 池化
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        
        # 第二个卷积块
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        
        # 第三个卷积块
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        
        # 展平特征图
        x = x.view(-1, 128 * 4 * 4)
        
        # 全连接层
        x = F.relu(self.fc1(x))
        x = self.dropout(x)  # 防止过拟合
        x = self.fc2(x)
        
        return x

# 创建模型实例并移动到设备
net = Net().to(device)

# 打印模型结构
print(net)

输出解释：

输出会显示我们定义的 CNN 模型的完整结构
模型包含 3 个卷积块（每个包含卷积层、批归一化层、ReLU 激活函数和池化层）
以及 2 个全连接层，最后输出 10 个类别的得分
我们使用了批归一化和 dropout 技术来提高模型性能并防止过拟合

步骤 5：定义损失函数和优化器

定义损失函数和优化器

# 定义损失函数（交叉熵损失）
criterion = nn.CrossEntropyLoss()

# 定义优化器（Adam）
optimizer = optim.Adam(net.parameters(), lr=0.001)

# 打印损失函数和优化器信息
print(f"损失函数: {criterion}")
print(f"优化器: {optimizer}")

输出解释：

我们使用交叉熵损失作为分类任务的损失函数
使用 Adam 优化器，学习率设置为 0.001
Adam 是一种常用的优化器，结合了动量和自适应学习率的优点

步骤 6：训练模型

现在我们开始训练模型，并观察训练过程中的损失和准确率变化。

# 训练模型
epochs = 15
train_losses = []
train_accs = []

for epoch in range(epochs):
    running_loss = 0.0
    correct = 0
    total = 0
    
    # 切换模型到训练模式
    net.train()
    
    for i, data in enumerate(train_loader, 0):
        # 获取输入数据
        inputs, labels = data[0].to(device), data[1].to(device)
        
        # 清零梯度
        optimizer.zero_grad()
        
        # 前向传播、计算损失、反向传播、参数更新
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        # 统计训练数据
        running_loss += loss.item()
        
        # 计算准确率
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        
        # 每100个批次打印一次信息
        if i % 100 == 99:
            print(f'[{epoch + 1}, {i + 1}] loss: {running_loss / 100:.3f}')
            running_loss = 0.0
    
    # 计算每个epoch的损失和准确率
    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100 * correct / total
    train_losses.append(epoch_loss)
    train_accs.append(epoch_acc)
    
    print(f'Epoch {epoch+1} - 损失: {epoch_loss:.3f}, 准确率: {epoch_acc:.2f}%')

print('Finished Training')

# 保存训练好的模型
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

# 绘制训练损失曲线
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(range(1, epochs+1), train_losses)
plt.title('Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')

# 绘制训练准确率曲线
plt.subplot(1, 2, 2)
plt.plot(range(1, epochs+1), train_accs)
plt.title('Training Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.tight_layout()
plt.show()

输出解释：

训练过程中会每 100 个批次打印一次当前的损失值
每个 epoch 结束后会打印该 epoch 的平均损失和准确率
训练完成后会保存模型参数到 cifar_net.pth 文件
最后会显示两张图表：训练损失曲线和训练准确率曲线

正常输出表现：

损失值应该随着训练的进行逐渐降低
准确率应该随着训练的进行逐渐提高
经过 15 个 epoch 后，训练准确率应该能达到 75-85% 左右

步骤 7：在测试集上评估模型

训练完成后，我们需要在测试集上评估模型的性能。

# 在测试集上评估模型
net.load_state_dict(torch.load(PATH))  # 加载保存的模型

correct = 0
total = 0
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))

# 切换模型到评估模式
net.eval()

# 不计算梯度，节省内存和计算资源
with torch.no_grad():
    for data in test_loader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        
        # 按类别统计准确率
        c = (predicted == labels).squeeze()
        for i in range(len(labels)):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1

# 打印总体准确率
print(f'测试集上的总体准确率: {100 * correct / total:.2f}%')

# 打印每个类别的准确率
for i in range(10):
    print(f'类别 {classes[i]} 的准确率: {100 * class_correct[i] / class_total[i]:.2f}%')

输出解释：

首先会显示模型在整个测试集上的总体准确率
然后会显示每个类别的准确率
对于 CIFAR-10 数据集，这个模型通常能达到 70-80% 的测试准确率
不同类别的准确率可能有差异，某些类别（如飞机、汽车）可能更容易识别，而其他类别（如猫、狗）可能更难识别

步骤 8：可视化模型预测结果

最后，让我们可视化一些测试图像和模型的预测结果。

# 展示一些测试图像和预测结果
dataiter = iter(test_loader)
images, labels = next(dataiter)

# 打印图像
imshow(torchvision.utils.make_grid(images[:5]))
print('真实标签: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(5)))

# 进行预测
outputs = net(images.to(device))
_, predicted = torch.max(outputs, 1)

print('预测标签: ', ' '.join(f'{classes[predicted[j]]:5s}' for j in range(5)))

# 展示一些错误预测的例子
misclassified_images = []
misclassified_labels = []
true_labels = []

net.eval()
with torch.no_grad():
    for data in test_loader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        
        # 找到错误预测的样本
        for i in range(len(labels)):
            if predicted[i] != labels[i]:
                misclassified_images.append(images[i].cpu())
                misclassified_labels.append(predicted[i].cpu())
                true_labels.append(labels[i].cpu())
                if len(misclassified_images) >= 5:
                    break
        if len(misclassified_images) >= 5:
            break

# 显示错误分类的图像
imshow(torchvision.utils.make_grid(misclassified_images[:5]))
print('真实标签: ', ' '.join(f'{classes[true_labels[j]]:5s}' for j in range(5)))
print('错误预测标签: ', ' '.join(f'{classes[misclassified_labels[j]]:5s}' for j in range(5)))

输出解释：