PyTorch 基于 ResNet18 的图像分类训练与验证全流程解析
一、项目概述
本文实现了一个基于 PyTorch 框架的图像分类模型，使用 ResNet18 作为预训练骨干网络（Backbone），并在其基础上进行迁移学习（Transfer Learning）。整个流程涵盖了：
-
数据预处理与增强
-
自定义 Dataset 与 DataLoader
-
模型微调与参数冻结
-
训练与验证循环
-
学习率调度策略（ReduceLROnPlateau）
该项目的核心目标是利用已有的强大视觉特征提取网络（ResNet18）对新的小规模数据集进行分类任务，从而提升训练效率与模型性能。
二、模型部分解析：ResNet18 微调机制
import torch
import torchvision.models as models
from torch import nn, optim
resnet_model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)这里加载了 torchvision.models 中的预训练 ResNet18 模型，其权重参数来自 ImageNet 大规模数据集的训练结果。
接着，冻结网络的所有参数，防止在训练过程中被更新：
for param in resnet_model.parameters():
print(param)
param.requires_grad = False✅ 原理说明：
迁移学习的关键思路在于"保留特征提取层"。
早期卷积层学习的是通用特征（如边缘、纹理），可直接用于新任务。
仅需微调后几层或分类头层（fc层），可显著减少训练量。
然后替换最后一层全连接层（fc层）：
in_features = resnet_model.fc.in_features
resnet_model.fc = nn.Linear(in_features, 20)🔹 原始 ResNet18 的输出为 1000 类（ImageNet）。
🔹 这里改为 20 类，适配自定义数据集。
最后仅选择需要更新的参数：
params_to_update = []
for param in resnet_model.parameters():
if param.requires_grad == True:
params_to_update.append(param)这意味着优化器只会更新新加入的全连接层参数。
三、数据预处理与增强（Data Augmentation）
数据增强可提升模型泛化能力，代码中定义了两种处理策略：
from torchvision import transforms1. 训练集增强 data_transforms['train']
包含大量随机性增强操作：
transforms.Compose([
transforms.Resize([300,300]),
transforms.RandomRotation(45),
transforms.CenterCrop(224),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomVerticalFlip(p=0.5),
transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),
transforms.RandomGrayscale(p=0.1),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])📘 数据增强效果：
旋转 、翻转 、灰度转换可提升模型在多视角条件下的鲁棒性。
归一化操作确保输入分布与预训练模型保持一致。
2. 验证集预处理 data_transforms['valid']
transforms.Compose([
transforms.Resize([256,256]),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])验证集通常不使用随机增强，以保持结果的可重复性和客观性。
四、自定义 Dataset 与 DataLoader
1. Dataset 类定义
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np
class food_dataset(Dataset):
def __init__(self, file_path, transform=None):
self.file_path = file_path
self.transform = transform
self.imgs = []
self.labels = []
with open(self.file_path) as f:
samples = [x.strip().split(' ') for x in f.readlines()]
for img_path,label in samples:
self.imgs.append(img_path)
self.labels.append(label)该类通过文本文件（如 train2.txt）加载图片路径和标签。
每一行格式为：
image_path label2. 数据访问接口
def __len__(self):
return len(self.imgs)
def __getitem__(self, index):
image = Image.open(self.imgs[index])
if self.transform:
image = self.transform(image)
label = torch.from_numpy(np.array(self.labels[index], dtype=np.int64))
return image, label⚙️ Dataset 必备方法：
__len__()：返回数据集大小。
__getitem__()：返回一条样本及其标签。
3. 数据加载器 DataLoader
training_data = food_dataset('./train2.txt', transform=data_transforms['train'])
test_data = food_dataset('./test2.txt', transform=data_transforms['valid'])
train_dataloader = DataLoader(training_data, batch_size=16, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=16, shuffle=True)✅ DataLoader 的作用：
自动打包 batch
支持多线程加载（
num_workers）支持数据打乱（shuffle）
五、训练环境与优化器设置
1. 自动选择设备
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")同时兼容：
-
NVIDIA GPU (CUDA)
-
Apple M1/M2 GPU (MPS)
-
CPU
2. 定义损失函数与优化器
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params_to_update, lr=0.001)使用 交叉熵损失函数 处理多分类任务，优化器为 Adam。
3. 学习率调度器
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode='min',
factor=0.5,
patience=3,
verbose=True
)🔹 当验证集 Loss 连续 3 轮未改善时，学习率减半。
🔹 可有效防止过拟合与梯度振荡。
六、训练与验证循环
1. 训练函数
def train(dataloader, model, loss_fn, optimizer):
model.train()
batch_size_num = 1
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()
if batch_size_num % 100 == 0:
print(f"loss: {loss.item():>7f} [number:{batch_size_num}]")
batch_size_num += 1🔁 每 100 个 batch 打印一次 loss。
2. 测试函数
def Test(dataloader, model, loss_fn):
size = len(dataloader.dataset)
num_batches = len(dataloader)
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 /= num_batches
correct /= size
print(f"Test result:\n Accuracy:{100*correct}%, Avg loss:{test_loss}")
return test_loss🧮 评估指标：
Accuracy（准确率）
Avg loss（平均验证损失）
七、主训练循环
epochs = 50
for t in range(epochs):
print(f"---------------\nepoch {t+1}")
train(train_dataloader, model, loss_fn, optimizer)
val_loss = Test(test_dataloader, model, loss_fn)
scheduler.step(val_loss)
print("Done!")共进行 50 轮训练，每轮包括：
-
模型训练
-
验证集测试
-
根据验证集 loss 调整学习率
💡 随着 epoch 增加，loss 应逐渐下降，accuracy 提升。
八、总结
| 模块 | 作用 | 特点 |
|---|---|---|
| ResNet18 | 特征提取主干 | 使用 ImageNet 预训练权重 |
| Dataset | 读取图片与标签 | 支持 transform 自动增强 |
| DataLoader | 批量化输入 | shuffle 提升训练效果 |
| train() | 前向传播与反向传播 | 更新梯度 |
| Test() | 模型评估 | 计算平均损失与准确率 |
| ReduceLROnPlateau | 学习率调整 | 自动降低学习率防止过拟合 |