- 🍨 本文为🔗365天深度学习训练营 中的学习记录博客
- 🍖 原作者：K同学啊
前言
本文记录一个完整的深度学习图像分类实践：使用 PyTorch 和迁移学习，将马铃薯叶片图像分为早疫病（Early blight）、晚疫病（Late blight）和健康（healthy）三类。项目同时实现了手写 VGG-16、预训练模型微调、训练评估、曲线可视化与单张图片预测。
目录
- 1. 项目简介
- 2. 项目环境
- 3. 数据集介绍
- 4. 数据预处理
- 5. 构建数据加载器
- 6. 手动搭建 VGG-16
- 7. 使用预训练 VGG-16 迁移学习
- 8. 模型训练与测试
- 9. 训练配置与模型保存
- 10. 结果可视化
- 11. 单张图片预测
1. 项目简介
马铃薯早疫病和晚疫病都会在叶片上形成明显病斑。传统人工识别依赖经验，面对大量作物时效率较低，因此可以使用卷积神经网络自动学习叶片颜色、纹理和病斑形态。
本项目的目标是输入一张马铃薯叶片图片，输出以下三种类别之一：
| 类别目录 | 中文含义 | 图片数量 |
|---|---|---|
Early_blight | 早疫病 | 1000 |
Late_blight | 晚疫病 | 1000 |
healthy | 健康叶片 | 152 |
| 合计 | 2152 |
项目采用 VGG-16 作为主干网络。为了理解网络结构，代码中先手动实现了一遍 VGG-16；正式训练时则加载 ImageNet 预训练权重，通过冻结部分参数和替换分类层完成迁移学习。
项目目录如下：
Potato_Disease_Identification/
├── Potato_Disease_Identification.py
└── PotatoPlants/
├── Early_blight/
├── Late_blight/
└── healthy/
2. 项目环境
项目代码依赖以下 Python 库：
Python
PyTorch
torchvision
Pillow
Matplotlib
torchsummary
可使用下面的命令安装主要依赖：
pip install torch torchvision pillow matplotlib torchsummary
程序会自动检测 CUDA。如果存在可用的 NVIDIA GPU，则使用 GPU 训练；否则退回 CPU。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
3. 数据集介绍
项目将数据集根目录转换为 Path 对象，并读取其中的类别文件夹：
import pathlib
data_dir = pathlib.Path("./PotatoPlants/")
data_paths = list(data_dir.glob("*"))
torchvision.datasets.ImageFolder 会把每个子目录视为一个类别，并按目录名排序生成标签映射。在当前数据集中，映射应为：
{
"Early_blight": 0,
"Late_blight": 1,
"healthy": 2
}
加载数据集：
total_data = datasets.ImageFolder(
"./PotatoPlants",
transform=train_transforms
)
print(total_data)
print(total_data.class_to_idx)
从数量上可以看出，健康叶片只有 152 张，而两个病害类别各有 1000 张，数据存在明显的类别不均衡。这一点可能导致模型更倾向于预测病害类别，后文会给出改进思路。
4. 数据预处理
VGG-16 的标准输入大小为 224 × 224。项目首先统一图片尺寸，再将图片转换为张量，最后使用 ImageNet 数据集的均值和标准差进行归一化：
from torchvision import transforms
train_transforms = transforms.Compose([
transforms.Resize([224, 224]),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
test_transforms = transforms.Compose([
transforms.Resize([224, 224]),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
归一化可以让输入分布与预训练 VGG-16 使用的 ImageNet 数据保持一致，有利于复用预训练特征。
当前训练集和测试集采用相同处理，没有加入数据增强。若希望提升泛化能力，可以在训练阶段增加随机翻转、轻微旋转和颜色扰动，但测试阶段应保持确定性的预处理。
5. 构建数据加载器
项目按照 8:2 划分训练集和测试集：
train_size = int(0.8 * len(total_data))
test_size = len(total_data) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(
total_data,
[train_size, test_size]
)
2152 张图片经过计算后得到：
| 数据部分 | 数量 |
|---|---|
| 训练集 | 1721 |
| 测试集 | 431 |
之后使用 DataLoader 按批次读取数据：
batch_size = 32
train_dl = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1
)
test_dl = DataLoader(
test_dataset,
batch_size=batch_size,
num_workers=1
)
读取一个测试批次时，典型张量形状为：
图像张量：torch.Size([32, 3, 224, 224])
标签张量：torch.Size([32])
四个维度分别表示批量大小、RGB 通道数、图片高度和图片宽度。
6. 手动搭建 VGG-16
VGG-16 的主要特点是连续使用 3 × 3 卷积核，并通过最大池化逐步缩小特征图。其卷积层按 2、2、3、3、3 分为五个模块，最后连接三层全连接分类器。
6.1 卷积模块
以第一个模块为例：
self.block1 = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(True),
nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(True),
nn.MaxPool2d(kernel_size=2, stride=2)
)
五次池化会把 224 × 224 的输入逐步变为 7 × 7 的特征图：
224 × 224
↓ block1
112 × 112
↓ block2
56 × 56
↓ block3
28 × 28
↓ block4
14 × 14
↓ block5
7 × 7
最后一个卷积模块输出 512 个通道，因此展平后的特征数为 512 × 7 × 7。
6.2 分类器
self.classifier = nn.Sequential(
nn.Linear(512 * 7 * 7, 4096),
nn.ReLU(),
nn.Linear(4096, 4096),
nn.ReLU(),
nn.Linear(4096, 3)
)
最后一层输出 3 个 logits，分别对应三个叶片类别。由于训练采用 CrossEntropyLoss，模型末尾不需要手动添加 Softmax。
6.3 前向传播
def forward(self, x):
x = self.block1(x)
x = self.block2(x)
x = self.block3(x)
x = self.block4(x)
x = self.block5(x)
x = torch.flatten(x, start_dim=1)
x = self.classifier(x)
return x
手动实现网络有助于理解 VGG-16，但从零训练需要更多数据和计算资源。因此项目正式训练时采用预训练模型。
7. 使用预训练 VGG-16 迁移学习
加载 torchvision 提供的 VGG-16：
vgg = torchvision.models.vgg16(pretrained=True).to(device)
首先冻结全部参数：
for param in vgg.parameters():
param.requires_grad = False
然后解冻后半部分卷积特征层：
for param in vgg.features[10:].parameters():
param.requires_grad = True
这样既可以保留预训练网络学到的通用低层特征，也允许较深层特征适应马铃薯叶片病斑。
原始 VGG-16 面向 ImageNet 的 1000 类分类任务，因此还需要将最后一层替换成 3 类输出：
vgg.classifier._modules["6"] = nn.Linear(
4096,
len(classNames)
)
vgg.to(device)
在较新版本的 torchvision 中，更推荐显式指定权重：
from torchvision.models import vgg16, VGG16_Weights
vgg = vgg16(weights=VGG16_Weights.DEFAULT).to(device)
8. 模型训练与测试
8.1 单轮训练
训练阶段包括前向传播、计算损失、梯度清零、反向传播和参数更新：
def train(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset)
num_batches = len(dataloader)
train_acc, train_loss = 0, 0
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()
train_acc += (pred.argmax(1) == y).float().sum().item()
train_loss += loss.item()
train_acc /= size
train_loss /= num_batches
return train_acc, train_loss
8.2 模型测试
测试阶段不需要计算梯度，使用 torch.no_grad() 可以减少显存占用并提高速度：
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset)
num_batches = len(dataloader)
test_acc, test_loss = 0, 0
with torch.no_grad():
for images, targets in dataloader:
images = images.to(device)
targets = targets.to(device)
pred = model(images)
loss = loss_fn(pred, targets)
test_acc += (pred.argmax(1) == targets).float().sum().item()
test_loss += loss.item()
test_acc /= size
test_loss /= num_batches
return test_acc, test_loss
8.3 完整训练循环
每个 epoch 完成后，程序记录训练集与测试集的准确率和损失，并输出当前学习率：
def Train(model, train_dl, test_dl, loss_fn, optimizer, epochs, schedule):
global train_acc, train_loss, test_acc, test_loss, best_acc
for epoch in range(epochs):
model.train()
epoch_train_acc, epoch_train_loss = train(
train_dl, model, loss_fn, optimizer
)
schedule.step()
model.eval()
epoch_test_acc, epoch_test_loss = test(
test_dl, model, loss_fn
)
if epoch_test_acc > best_acc:
best_acc = epoch_test_acc
best_model = copy.deepcopy(model)
train_acc.append(epoch_train_acc)
train_loss.append(epoch_train_loss)
test_acc.append(epoch_test_acc)
test_loss.append(epoch_test_loss)
lr = optimizer.state_dict()["param_groups"][0]["lr"]
print(
f"Epoch:{epoch + 1:2d}, "
f"Train_acc:{epoch_train_acc * 100:.1f}%, "
f"Train_loss:{epoch_train_loss:.3f}, "
f"Test_acc:{epoch_test_acc * 100:.1f}%, "
f"Test_loss:{epoch_test_loss:.3f}, "
f"Lr:{lr:.2E}"
)
return best_model
代码通过深拷贝保留测试准确率最高的模型，避免只使用最后一个 epoch 的参数。
9. 训练配置与模型保存
项目采用交叉熵损失、Adam 优化器和 StepLR 学习率调度器：
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(
vgg.parameters(),
lr=1e-6
)
schedule = torch.optim.lr_scheduler.StepLR(
optimizer=optimizer,
gamma=0.9,
step_size=15
)
epochs = 60
参数含义如下：
| 参数 | 值 | 作用 |
|---|---|---|
batch_size | 32 | 每次输入模型的图片数量 |
epochs | 60 | 完整遍历训练集的次数 |
lr | 1e-6 | Adam 初始学习率 |
step_size | 15 | 每 15 个 epoch 调整一次学习率 |
gamma | 0.9 | 学习率乘以 0.9 |
启动训练并保存权重：
best_model = Train(
vgg,
train_dl,
test_dl,
loss_fn,
optimizer,
epochs,
schedule
)
torch.save(vgg.state_dict(), "teacher_vgg16_best.pth")
需要注意：上面的保存语句保存的是 vgg 当前状态，也就是最后一个 epoch，而不是变量 best_model 中的最佳状态。更准确的写法是：
torch.save(
best_model.state_dict(),
"teacher_vgg16_best.pth"
)
10. 结果可视化
项目使用 Matplotlib 绘制准确率和损失曲线：
def Results_visualization(epochs):
import matplotlib.pyplot as plt
epochs_range = range(epochs)
plt.figure(figsize=(12, 3))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_acc, label="Training Accuracy")
plt.plot(epochs_range, test_acc, 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, label="Training Loss")
plt.plot(epochs_range, test_loss, label="Test Loss")
plt.legend(loc="upper right")
plt.title("Training and Validation Loss")
plt.show()
11. 单张图片预测
训练完成后，可以读取一张图片并完成预测：
def predict_one_image(image_path, model, transform, classes):
test_img = Image.open(image_path).convert("RGB")
test_img = transform(test_img)
image = test_img.to(device).unsqueeze(0)
model.eval()
with torch.no_grad():
output = model(image)
probability = torch.softmax(output, dim=1)
confidence, pred = torch.max(probability, 1)
pred_class = classes[pred.item()]
print(f"预测结果：{pred_class}")
print(f"置信度：{confidence.item() * 100:.2f}%")
调用示例：
classes = total_data.classes
predict_one_image(
"./PotatoPlants/Late_blight/example.JPG",
best_model,
test_transforms,
classes
)
这里使用 unsqueeze(0) 增加批次维度，把单张图片由 [3, 224, 224] 变成 [1, 3, 224, 224]。同时使用 softmax 将 logits 转换为概率，以便输出预测置信度。
