之前使用 TensorFlow 图像分类，现在使用 PyTorch 进行重现，实现相同的效果：
- CNN 识别图片
- 从磁盘高效加载数据集。
- 识别过拟合，并应用数据增强和随机失活等技术缓解过拟合。
1、导入 PyTorch 和其他必要的库
使用 pip 安装相关依赖库
pip install torch torchvision matplotlib numpy==1.26.4
代码中引入相关依赖
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split
import matplotlib.pyplot as plt
import numpy as np
2、数据集加载
使用与 TensorFlow 图像分类 相同的数据集，包含约 3,700 张花卉照片的数据集。该数据集包含 5 个子目录，每个子目录对应一个类。相关图片说明也在之前文章里解释过，不再赘述。
这里从 Google API 下载到本地，然后从本地加载，并进行水平翻转、旋转进行数据增强。然后把数据集分为训练集和验证集，80% 的图像用于训练，将 20% 的图像用于验证。
# ==================== 1. 超参数配置 ====================
data_dir = './flower_photos' # 图片根目录
train_ratio = 0.8 # 80%训练，20%验证
num_epochs = 10 # 训练批次
batch_size = 32
learning_rate = 0.001
# ==================== 2. 数据预处理，包括数据增强 ====================
train_transform = transforms.Compose([
transforms.Resize((64, 64)), # 统一缩放到64x64，加快训练
transforms.RandomHorizontalFlip(), # 随机水平翻转（简单增强）
transforms.RandomRotation(15),
transforms.ToTensor(), # 转换为Tensor [0,1]
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) # 归一化到[-1,1]
])
val_transform = transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
# ==================== 3. 加载并分割数据集 ====================
print("正在加载并分割数据集...")
# 加载完整数据集
full_dataset = datasets.ImageFolder(root=data_dir, transform=train_transform)
# 计算训练集和验证集的大小
dataset_size = len(full_dataset)
train_size = int(train_ratio * dataset_size)
val_size = dataset_size - train_size
# 随机分割数据集
train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])
# 为验证集单独设置转换
val_dataset.dataset.transform = val_transform
print(f'总图片数: {dataset_size} 张')
print(f'训练集: {train_size} 张, 验证集: {val_size} 张')
print(f'类别数: {len(full_dataset.classes)} 个')
print(f'类别对应关系: {full_dataset.class_to_idx}')
# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
3、定义模型
整体流程为：输入图像 → 三个卷积层 → 展平 → 全连接层 → 类别预测，这个模型结构虽简单但完整，涵盖了 CNN 的核心组件：
- 卷积层 (
nn.Conv2d)：学习局部特征 - 激活函数 (
nn.ReLU)：引入非线性 - 池化层 (
nn.MaxPool2d)：下采样，减少计算量并增加感受野 - 全连接层 (
nn.Linear)：综合所有特征进行分类 -
Dropout：随机丢弃部分神经元，防止过拟合
# ==================== 3. 定义模型 ====================
class SimpleFlowerCNN(nn.Module): # 继承自nn.Module，在PyTorch中，所有神经网络模块都继承自nn.Module。
def __init__(self, num_classes):
super(SimpleFlowerCNN, self).__init__()
# 卷积块1: 输入3通道(彩色), 输出16通道
self.conv1 = nn.Sequential(
nn.Conv2d(3, 16, kernel_size=3, padding=1), # 输出: (16, 64, 64)
nn.ReLU(),
nn.MaxPool2d(2) # 输出: (16, 32, 32)
)
# 卷积块2: 输入16通道, 输出32通道
self.conv2 = nn.Sequential(
nn.Conv2d(16, 32, kernel_size=3, padding=1), # 输出: (32, 32, 32)
nn.ReLU(),
nn.MaxPool2d(2) # 输出: (32, 16, 16)
)
# 卷积块3: 输入32通道, 输出64通道
self.conv3 = nn.Sequential(
nn.Conv2d(32, 64, kernel_size=3, padding=1), # 输出: (64, 16, 16)
nn.ReLU(),
nn.MaxPool2d(2) # 输出: (64, 8, 8)
)
# 全连接层
self.fc = nn.Sequential(
nn.Flatten(), # 将64*8*8=4096维特征展平
nn.Linear(64 * 8 * 8, 128), # 第一全连接层
nn.ReLU(),
nn.Dropout(0.5), # 随机失活，防止过拟合
nn.Linear(128, num_classes) # 输出层，对应类别数
)
# 前向传播函数定义了数据如何通过模型。
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = self.fc(x)
return x
其中，卷积层1的定义参数含义：
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
-
in_channels：输入通道数 3，因为输入图像是彩色 RGB 三通道，所以是 3，如果是灰度图，则为1。 -
out_channels：输出通道数 16，表示使用 16 个不同的卷积核来提取 16 种特征。这个数字通常是 2 的幂次（如 16、32、64），但也可以随意设定。一般越深的层，通道数越多，以学习更复杂的特征。 -
kernel_size：卷积核的大小是 3×3。常见的有 3×3、5×5 等。3×3 是最常用的，因为它能捕捉局部特征且参数较少。 -
padding：在输入图像的四周填充1圈0。
4、设置训练参数
初始化模型，并定义损失函 nn.CrossEntropyLoss() 交叉熵损失，适用于多分类任务，以及优化器 AdamW，传入模型中所有需要训练的参数，并设置学习率为 0.001。
# ==================== 4. 训练设置 ====================
# 初始化模型
num_classes = len(full_dataset.classes)
model = SimpleFlowerCNN(num_classes)
print(f"\n模型类别数: {num_classes}")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
print(f"使用设备: {device}")
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=learning_rate)
5、开始训练
训练模型的代码基本是固定的流程：
- 把输入特征给到模型，得到模型的预测值；
- 把预测值和标签给到损失函数，计算损失值；
- 损失值反响传播，计算梯度；
- 根据梯度更新模型参数。
其中增加了一些统计模型准确度、损失值的计算，这部分仅仅是画图统计使用，可以忽略。
# ==================== 5. 训练循环（记录指标） ====================
print("\n开始训练...")
# 创建列表来记录训练过程中的指标
train_losses = []
train_accuracies = []
val_losses = []
val_accuracies = []
for epoch in range(num_epochs):
# 训练阶段
model.train()
train_loss, train_correct, train_total = 0.0, 0, 0
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device) # 数据移动到设备
optimizer.zero_grad() # 清空梯度，避免累加
outputs = model(inputs) # 模型预测
loss = criterion(outputs, labels) # 计算损失
loss.backward() # 反向传播，计算梯度
optimizer.step() # 根据梯度更新参数
train_loss += loss.item()
_, predicted = outputs.max(1)
train_total += labels.size(0)
train_correct += predicted.eq(labels).sum().item()
# 计算训练指标
avg_train_loss = train_loss / len(train_loader)
train_acc = 100. * train_correct / train_total
train_losses.append(avg_train_loss)
train_accuracies.append(train_acc)
# 验证阶段
model.eval()
val_loss, val_correct, val_total = 0.0, 0, 0
with torch.no_grad(): # 验证时不计算梯度
for inputs, labels in val_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = outputs.max(1)
val_total += labels.size(0)
val_correct += predicted.eq(labels).sum().item()
# 计算验证指标
avg_val_loss = val_loss / len(val_loader)
val_acc = 100. * val_correct / val_total
val_losses.append(avg_val_loss)
val_accuracies.append(val_acc)
print(f'Epoch [{epoch + 1:02d}/{num_epochs}] | '
f'Train Loss: {avg_train_loss:.4f}, Acc: {train_acc:.2f}% | '
f'Val Loss: {avg_val_loss:.4f}, Acc: {val_acc:.2f}%')
print("训练完成！")
# 保存模型
torch.save(model.state_dict(), 'simple_flower_cnn.pth')
print("模型已保存为 'simple_flower_cnn.pth'")
重要说明：
-
loss.backward(): 计算损失相对于每个参数的梯度 -
optimizer.step(): 使用梯度更新参数（参数 = 参数 - lr × 梯度） -
optimizer.zero_grad(): 必须清零梯度，否则梯度会累加
最后画图统计结果
可以看到损失值随着训练批次的增加逐渐减小，准确率随着训练批次增加逐渐提高，这是符合我们预期的结果。最终准确率在 70% 左右，主要原因还是训练数据太少。
6、迁移学习
实际上训练数据只有约3000张图片，训练数据太少，搭建一个全新的 cnn 网络并不是一个好方法，容易导致过拟合（模型只记住了训练集，而无法泛化到新图片）。因此，最佳策略是迁移学习。这里以经典的 ResNet-50 为例。
ResNet-50 作为深度残差网络的经典代表，其核心突破在于引入残差连接（Residual Connection）机制。传统深度神经网络面临梯度消失或爆炸问题，导致深层网络训练困难。ResNet通过”捷径连接”（Shortcut Connection）将输入直接传递到深层，形成恒等映射（Identity Mapping），使得网络可以专注于学习残差部分（F(x)=H(x)-x），从而有效缓解梯度消失问题。
具体架构上，ResNet-50 包含 49 个卷积层和 1 个全连接层，总参数量约 2550 万。其核心模块为 Bottleneck 结构，由 1×1、3×3、1×1 三个卷积层组成：第一个1×1卷积用于降维（减少计算量），3×3卷积提取特征，第二个1×1卷积恢复维度。这种设计在保持特征表达能力的同时，将计算复杂度从标准残差块的O(k²)降至O(k)，其中k为卷积核尺寸。
要把上面我们自定义的 cnn 转为迁移学习，只需要 3 步：
- 数据预处理
- 加载模型
- 修改优化器
1、数据预处理
ResNet-50 输入图像尺寸需调整为 224×224（ResNet-50标准输入），预训练模型时必须采用相同的标准化参数，而之前是把图片缩放到 64×64，因此数据预处理部分需要修改。
# ==================== 2. 数据预处理 ====================
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224), # 随机裁剪并缩放至224x224
transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.RandomRotation(20), # 随机旋转
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # ImageNet标准归一化
])
val_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])
])
2、加载模型
其次需要修改的地方就是不再使用我们自定义的 cnn 模型，而是加载 ResNet-50 模型，并冻结权重，只有最后的全连接层进行替换，替换为我们花卉分类所需要的。
# 3. 加载预训练模型并微调
model = models.resnet50(pretrained=True) # 加载预训练权重
# 冻结所有模型参数（只微调最后的全连接层）
for param in model.parameters():
param.requires_grad = False
# 替换最后的全连接层，适配花卉类别数
# 替换最后的全连接层
num_features = model.fc.in_features
model.fc = nn.Sequential(
nn.Linear(num_features, 512),
nn.BatchNorm1d(512),
nn.ReLU(inplace=True),
nn.Dropout(p=0.4), # Dropout对抗过拟合
nn.Linear(512, num_classes)
)
3、修改优化器
最后修改优化器，只有最后的全连接层需要参数的优化传播，其他层都不需要。
optimizer = optim.AdamW(model.fc.parameters(), lr=learning_rate)
现在就可以重现训练模型了。直接看效果。
由于模型较大，训练时间较长，这里只是演示，因此只训练 5 次，可以看到，训练准确率直接到 90%，比之前增加了 20%，效果非常明显。
