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

一、前期准备

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets, transforms, models    #"借"大脑
from torch.utils.data import DataLoader, random_split   #"切"数据
import matplotlib.pyplot as plt
import numpy as np
import os
from PIL import Image

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

二、导入数据

# 我们还是用 ImageNet 的“魔法数字”
data_dir = './data' # 就是那个有 Monkeypox 文件夹的 data 文件夹

# “练习册”的“食谱”（带“增强”）
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# “考试卷”的“食谱”（干干净净）
test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 用 ImageFolder 去读那个 data 文件夹
# P4 的“食材”少，我们把“练习册”和“考试卷”用同一个“食谱”
# (我们用“考试卷”那个“干净”的，不然“增强”的“练习册”和“干净”的“考试卷”没法 random_split)
dataset = datasets.ImageFolder(data_dir, transform=test_transform)
print("成功读取 ImageFolder！")

# P4 也要自己“切”练习册和考试卷 (和 P3 一样)
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_ds, test_ds = random_split(dataset, [train_size, test_size])

print(f"总图片数: {len(dataset)}")
print(f"练习册(80%): {len(train_ds)} 张")
print(f"考试卷(20%): {len(test_ds)} 张")

# 找“喂饭小助手”
batch_size = 32
train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
test_dl = DataLoader(test_ds, batch_size=batch_size, shuffle=False)

# P4 的 2 个分类名字 (['Monkeypox', 'Others'])
classes = dataset.classes
print(f"分类: {classes}")

三、构建神经网络 (P4 的“迁移学习”)

# 去借一个“超级大脑” ResNet34
model = models.resnet34(pretrained=True)

# 把“大脑”的“身体”全都“冻住”
for param in model.parameters():
    param.requires_grad = False

# “砍掉”它原来的“头”(model.fc)，换上我们自己的“新头”
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 2)

# 把 P4“大脑”搬到GPU上
model = model.to(device)

四、训练模型

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

learning_rate = 0.001
epochs = 20 # 跑 20 轮

loss_fn = nn.CrossEntropyLoss()
# 告诉“辅导员”，只“辅导”那个“新头”！
optimizer = optim.Adam(model.fc.parameters(), lr=learning_rate)

4.2 编写训练和测试函数

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

4.3 正式训练（循环执行）

print("开始正式训练... (只训练那个'新头')")
train_loss_history = []
train_acc_history = []
test_loss_history = []
test_acc_history = []

for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}\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!")

五、保存模型

model_path = "P4_Monkeypox_ResNet34.pth"
torch.save(model.state_dict(), model_path)
print(f"模型已保存到: {model_path}")

六、结果可视化

6.1 Loss和Accuracy图

import warnings
warnings.filterwarnings("ignore")
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

epochs_range = range(epochs)
plt.figure(figsize=(12, 3))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_acc_history, label='Training Accuracy')
plt.plot(epochs_range, test_acc_history, label='Test Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')
plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_loss_history, label='Training Loss')
plt.plot(epochs_range, test_loss_history, label='Test Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()

6.2 预测

找一张“猴痘” 或者“其他” 的图片"test_pox.jpg"。

print("\n开始预测一张新图片...")

predict_transform = test_transform # P4 用 test_transform 就行

img_path = "./test_pox.jpg" 

try:
    img = Image.open(img_path).convert('RGB')
    img_tensor = predict_transform(img)
    img_tensor = img_tensor.unsqueeze(0)
    img_tensor = img_tensor.to(device)

    model.eval()
    with torch.no_grad():
        pred = model(img_tensor)
        
    probabilities = torch.softmax(pred, dim=1)
    probability, predicted_class_index = torch.max(probabilities, 1)
    predicted_class_name = classes[predicted_class_index.item()]
    
    print(f"\n--- 最终预测结果 ---")
    print(f"我猜这张图是: {predicted_class_name}")
    print(f"我有多自信 (概率): {probability.item() * 100:.2f}%")

    # 画图
    plt.figure()
    plt.imshow(img.resize((224, 224)))
    plt.title(f"预测: {predicted_class_name} ({probability.item()*100:.2f}%)")
    plt.axis('off')
    plt.show()

except FileNotFoundError:
    print(f"!!! 没找到 '{img_path}' !!!")

七、总结

P4注意点：

1. 迁移学习 (Transfer Learning)
这里引入了预训练模型 (Pre-trained Model)。
我们加载了 models.resnet34(pretrained=True)， 这意味着我们“借用”了一个已经在 ImageNet (一个巨大的数据集) 上训练好的“超级模型”。这能极大地提高训练速度和最终精度，尤其是在我们自己的“猴痘” 数据集很小的情况下。

2. 特征冻结 (Feature Freezing) 与分类头替换

P3 里我们训练（更新）模型中的所有参数，而在在P4这里我们只训练模型的最后一部分。
注意点：
冻结： 我们用 param.requires_grad = False 循环，“冻结”了 resnet34 的所有“卷积基”（身体）。这是因为它的“身体”已经学会了如何提取通用的图像特征（如边缘、纹理）。

替换： 我们必须替换掉它原来的“分类头”(model.fc)， 因为它原来是为 1000 个分类设计的。我们通过 model.fc = nn.Linear(num_features, 2) 将其替换为只适应我们 2 个分类（Monkeypox, Others）的新层。

3. 优化器的目标
在P4优化器只更新我们“新换上去的头”。
这是“迁移学习” 的关键一步。我们传给优化器的是 optim.Adam(model.fc.parameters(), …)， 明确告诉它，只有 model.fc（新头）的参数需要被“辅导”（更新）。

4. 数据处理流程 (Transforms)
P4的数据集在物理上就被分为了 train 和 test 两个文件夹。
这才是专业的做法。它允许我们定义两种不同的 transform：

train_transform：包含数据增强 (Data Augmentation)，如 RandomResizedCrop 和 RandomHorizontalFlip。 这能防止过拟合，提高模型的“泛化能力”。

test_transform：绝不能包含随机性。它必须是确定的（如 Resize 和 CenterCrop），以确保我们每次“考试” 都是在同一套标准下进行的。

5. 标准化 (Normalization)
这里使用 Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])。
这组“魔法数字” 是 ImageNet 数据集的均值和标准差。因为我们借的 resnet34 是用这组数字“喂”大的，所以我们必须用一模一样的“食谱” 去“加工”我们的新食材， 它才能认识。

6. 保存模型 (Save Model)
我们在训练结束后， 使用了 torch.save(model.state_dict(), …)。
这是为了把我们辛辛苦苦“只训练了‘新头’” 的模型保存下来。这样我们下次“预测” 时，就不需要再花 20 轮 的时间去重新训练了。