关注我天天好文
- 实现：用 PyTorch + torchvision + timm + ResNet-18 迁移学习完成多类别图像分类。
- 流程：检查数据目录 → 设置随机种子与设备 → 加载预训练模型及其预处理配置 → 构建数据集 → 训练与验证 → 保存最佳模型 → 在测试集上进行一次最终评估。
- 常见问题与排查
- CUDA 显存不足：减小
batch_size，例如从 32 改为 16 或 8。 - 首次运行下载权重失败：确认网络可访问模型权重地址，或提前将权重缓存到本地。
- 验证准确率波动：确认验证集没有随机增强，且训练、验证、测试没有样本泄漏。
- 训练集准确率高而验证集差：检查数据划分、类别不平衡、数据增强强度与过拟合；不能根据测试集调参。
- 预训练配置接口报错：检查
timm版本；本文使用timm.data.resolve_model_data_config接口。
- CUDA 显存不足：减小
数据目录要求
保证目录结构如下。每个类别必须是一个子目录，图片放在对应类别目录中。
train/
apple/
banana/
...
validation/
apple/
banana/
...
test/
apple/
banana/
...
- 三个数据集必须包含相同的类别目录。
ImageFolder会按类别目录名称的字典序分配标签，因此这里会检查三份数据集的类别映射是否一致。 - 首次创建预训练模型时，
timm可能需要下载模型权重；请确保环境已安装兼容版本的torch、torchvision和timm。
1. 导入依赖
torch：模型训练、自动求导和张量运算。torchvision：图像变换与ImageFolder数据集。timm：创建带预训练权重的 ResNet-18。matplotlib：显示样例图片。- 不再使用旧版
imgaug：这里用torchvision.transforms完成常见增强，减少额外依赖及 NumPy 兼容问题。
import os
import random
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import timm
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
2. 设置随机种子
随机种子用于控制 Python、NumPy 和 PyTorch 中常见的随机过程，使实验更容易复现。
- 固定随机种子不代表所有硬件、驱动和软件版本下都能得到完全相同的结果。
- 关闭 cuDNN benchmark 并启用 deterministic，有助于提高 CUDA 训练的可复现性，但某些算子可能因此变慢或不支持确定性实现。
def set_seed(seed: int = 42) -> None:
# Python 和 NumPy 的随机过程
random.seed(seed)
np.random.seed(seed)
# PyTorch CPU / CUDA 随机过程
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
# 优先保证可复现性，而不是自动搜索最快的卷积算法
if torch.backends.cudnn.is_available():
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
set_seed(42)
3. 配置数据路径、设备和超参数
设备选择顺序为 CUDA → Apple Silicon 的 MPS → CPU
batch_size：每次参数更新使用的样本数量；显存不足时可调小。lr：学习率。这里对所有模型参数使用同一个学习率，作为简单的迁移学习基线。epochs：完整遍历训练集的次数。num_workers=0：先采用最稳妥的数据加载设置，确认运行正常后再根据系统和数据加载速度调整。
# 数据集根目录：默认位于当前工作目录
train_path = Path("./train")
val_path = Path("./validation")
test_path = Path("./test")
# 先检查路径，避免后续报出不直观的 FileNotFoundError
for path in (train_path, val_path, test_path):
if not path.is_dir():
raise FileNotFoundError(
f"找不到数据目录：{path.resolve()}。请检查数据集是否已解压到当前工作目录。"
)
# 自动选择可用设备
if torch.cuda.is_available():
device = torch.device("cuda")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")
batch_size = 32
num_workers = 0
lr = 3e-4
epochs = 10
seed = 42
print(f"设备: {device}")
print(f"Batch size: {batch_size}")
print(f"学习率: {lr}")
print(f"训练轮数: {epochs}")
4. 创建预训练模型，并读取对应的预处理配置
使用 timm.create_model 加载 ImageNet 预训练的 ResNet-18，并将分类头的输出类别数设为数据集的实际类别数。
为什么先读取模型的数据配置？ 预训练权重通常对应特定的输入尺寸、通道均值和标准差。直接对所有模型硬编码一组归一化参数，可能与所选权重的训练预处理不一致。这里使用 timm 提供的数据配置接口构造评估变换。
pretrained=True：加载预训练权重。num_classes：替换分类头，使输出维度与当前任务匹配。resolve_model_data_config：获取模型推荐的数据预处理配置。
# 先加载训练集目录，以确定类别映射和类别数量
train_dataset_raw = ImageFolder(str(train_path))
categories = train_dataset_raw.classes
num_classes = len(categories)
if num_classes < 2:
raise ValueError(f"至少需要两个类别进行分类，当前检测到 {num_classes} 个类别。")
print(f"类别数量: {num_classes}")
print(f"类别名称: {categories}")
# 创建预训练模型；首次运行时可能需要下载权重
model = timm.create_model(
"resnet18",
pretrained=True,
num_classes=num_classes
).to(device)
# 根据当前模型的预训练配置生成输入预处理参数 读取预训练模型配置
data_config = timm.data.resolve_model_data_config(model)
input_size = data_config["input_size"]
input_height, input_width = input_size[1], input_size[2]
print(f"模型输入尺寸: {input_size}")
print(f"输入均值: {data_config['mean']}")
print(f"输入标准差: {data_config['std']}")
5. 构建训练与评估变换
训练集采用随机增强，提高模型对常见姿态和外观变化的适应能力；验证集和测试集只做确定性预处理，避免每次评估都因随机增强而波动。
RandomHorizontalFlip：随机水平翻转；适用于左右翻转不改变类别语义的任务。RandomRotation：小角度旋转，模拟拍摄角度变化。ColorJitter：轻微调整颜色，提升对光照和色彩变化的鲁棒性。ToTensor：将 PIL 图像转换为[C, H, W]的浮点张量，并将像素值缩放到[0, 1]。Normalize：使用模型预训练配置中的均值和标准差进行标准化。- 注意： 验证集和测试集不使用随机增强；显示已标准化的张量时，需要先反标准化。
train_tfs = transforms.Compose([
transforms.Resize((input_height, input_width)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomRotation(degrees=20),
transforms.ColorJitter(
brightness=0.1,
contrast=0.1,
saturation=0.1,
hue=0.03
),
transforms.ToTensor(),
transforms.Normalize(
mean=data_config["mean"],
std=data_config["std"]
),
])
eval_tfs = transforms.Compose([
transforms.Resize((input_height, input_width)),
transforms.ToTensor(),
transforms.Normalize(
mean=data_config["mean"],
std=data_config["std"]
),
])
6. 创建 Dataset 与 DataLoader
ImageFolder 根据子目录组织图像，并自动生成整数标签。训练集开启 shuffle=True，让每轮训练的样本顺序变化；验证集和测试集关闭打乱，便于稳定评估和结果追踪。
这里显式检查三个数据集的类别名称与标签映射是否一致，防止训练和评估阶段的同一整数标签代表不同类别。
train_dataset = ImageFolder(str(train_path), transform=train_tfs)
val_dataset = ImageFolder(str(val_path), transform=eval_tfs)
test_dataset = ImageFolder(str(test_path), transform=eval_tfs)
# 检查类别名称与索引映射一致
assert train_dataset.classes == val_dataset.classes, (
"训练集与验证集的类别目录不一致，或标签映射不同。"
)
assert train_dataset.classes == test_dataset.classes, (
"训练集与测试集的类别目录不一致，或标签映射不同。"
)
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=(device.type == "cuda")
)
val_loader = DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=(device.type == "cuda")
)
test_loader = DataLoader(
test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=(device.type == "cuda")
)
print(f"训练集: {len(train_dataset)} 张")
print(f"验证集: {len(val_dataset)} 张")
print(f"测试集: {len(test_dataset)} 张")
print(f"类别数: {len(train_dataset.classes)}")
7. 检查一个 Batch 的形状并显示样例
模型输入应为 [batch_size, channels, height, width]，例如 [32, 3, 224, 224]。
由于张量已经执行标准化，不能直接将其作为 RGB 像素显示。下面先反标准化，再把通道维从 CHW 转回 HWC 供 Matplotlib 显示。
images, labels = next(iter(train_loader))
print(f"图像 Batch 形状: {tuple(images.shape)}")
print(f"标签 Batch 形状: {tuple(labels.shape)}")
print(f"图像张量类型: {images.dtype}")
print(f"标签示例: {labels[:8].tolist()}")
# 反标准化：x_original = x_normalized * std + mean
mean = torch.tensor(data_config["mean"]).view(3, 1, 1)
std = torch.tensor(data_config["std"]).view(3, 1, 1)
image_to_show = (images[0].cpu() * std + mean).clamp(0, 1)
image_to_show = image_to_show.permute(1, 2, 0).numpy()#排布成(H, W, C)
plt.figure(figsize=(5, 5))
plt.imshow(image_to_show)
plt.title(f"Label: {categories[labels[0].item()]}")
plt.axis("off")
plt.show()
8. 查看 timm 模型库与当前模型配置
这部分是辅助探索，不参与训练过程。
timm.list_models(pretrained=True)：列出 timm 中可用预训练权重的模型名称。timm.list_models("*resnet*")：按通配符筛选 ResNet 相关模型。- 当前模型的数据配置已在前面通过
resolve_model_data_config读取；不再用旧式model.default_cfg写法。
pretrained_models = timm.list_models(pretrained=True)
print(f"可查询到的预训练模型数量: {len(pretrained_models)}")
resnet_models = timm.list_models("*resnet*")
print(f"匹配到的 ResNet 模型数量: {len(resnet_models)}")
print(resnet_models[:50]) # 仅展示前 50 个，避免输出过长
9. 定义损失函数、优化器与学习率调度器
CrossEntropyLoss：多分类任务的常用损失函数，输入是模型输出的 logits 和整数类别标签；不要模型输出后做 Softmax。Adam：自适应梯度优化器。CosineAnnealingLR：按余弦曲线逐步降低学习率。- 每个 Epoch 结束后调用一次
scheduler.step()，所以T_max=epochs表示以 Epoch 为单位的调度周期。如果改成每个 Batch 更新一次，就必须相应地按总 Batch 步数设置T_max。
from torch.optim import Adam
from torch.optim.lr_scheduler import CosineAnnealingLR
criterion = nn.CrossEntropyLoss()
# 定义 Adam 优化器，用于根据梯度更新模型的权重参数
optimizer = Adam(
model.parameters(),
lr=lr, # 初始学习率（Learning Rate）
betas=(0.9, 0.999), # 动量的一阶矩和二阶矩估计指数衰减率（默认值）
eps=1e-8, # 数值稳定常数，避免分母为 0
weight_decay=0.0 # 权重衰减（L2 正则化系数，此处未启用）
)
# 定义余弦退火学习率调度器（Cosine Annealing），让学习率随训练轮数呈余弦曲线平滑下降
scheduler = CosineAnnealingLR(
optimizer,
T_max=epochs, # 达到最小学习率所需的总训练轮数（半个余弦周期）
eta_min=1e-6 # 学习率衰减到的最低下限值
)
10. 训练一个 Epoch
训练步骤：前向传播 → 计算损失 → 清空旧梯度 → 反向传播 → 更新参数。
model.train()开启训练模式，使 BatchNorm、Dropout 等模块按训练模式工作。optimizer.zero_grad()清除上一步累积的梯度。loss.backward()计算梯度。optimizer.step()更新模型参数。- 统计损失时用每个 Batch 的平均损失乘以该 Batch 的样本数，最后除以数据集样本总数；这样即使最后一个 Batch 较小，平均损失仍然正确。
def train_one_epoch(epoch, model, train_loader):
model.train()
total_loss = 0.0
total_samples = 0
# 用 tqdm 包装训练数据加载器（train_loader），生成一个带动态进度条的迭代器
progress = tqdm(
train_loader,
desc=f"Train {epoch + 1}/{epochs}" # 进度条前缀文本，实时显示当前训练轮数（如 Train 1/10）
)
for images, targets in progress:
images = images.to(device)
targets = targets.to(device)
# 清除上一个 Batch 的梯度
optimizer.zero_grad(set_to_none=True)
# 前向传播：输出形状为 [batch_size, num_classes] 的 logits
logits = model(images)
loss = criterion(logits, targets)
# 反向传播与参数更新
loss.backward()
optimizer.step()
batch_samples = images.size(0)
total_loss += loss.item() * batch_samples
total_samples += batch_samples
# 在 tqdm 进度条的右侧动态追加/更新显示当前的损失值（loss），将 Tensor 标量转为 Python 浮点数并保留 4 位小数
progress.set_postfix(loss=f"{loss.item():.4f}")
return total_loss / total_samples
11. 在验证集上评估
验证阶段不更新参数。
model.eval()切换到评估模式。torch.inference_mode()关闭梯度记录，适合纯推理与评估，通常比普通梯度模式开销更低。- 验证损失按样本总数求平均，避免错误地除以 Batch 数。
- 验证集用于模型选择和训练过程监控；测试集不应在每个 Epoch 都参与评估。
def evaluate(model, data_loader, split_name="Validation"):
model.eval()
total_loss = 0.0
total_samples = 0
correct = 0
all_targets = []
all_preds = []
with torch.inference_mode():
for images, targets in tqdm(data_loader, desc=split_name):
images = images.to(device)
targets = targets.to(device)
logits = model(images)
loss = criterion(logits, targets)
preds = logits.argmax(dim=1)
batch_samples = images.size(0)
total_loss += loss.item() * batch_samples
total_samples += batch_samples
correct += (preds == targets).sum().item()
all_targets.append(targets.cpu().numpy())
all_preds.append(preds.cpu().numpy())
avg_loss = total_loss / total_samples
accuracy = correct / total_samples
return {
"loss": avg_loss,
"accuracy": accuracy,
"targets": np.concatenate(all_targets),
# 沿着默认轴（axis=0）拼接融合成一个完整的 NumPy 数组，通常作为字典的值用于后续计算整体评估指标
"preds": np.concatenate(all_preds),
}
12. 训练循环：使用验证集选择最佳模型
每个 Epoch 完成训练后，在验证集上评估；当验证准确率提升时保存模型参数。
- 保存
state_dict，即模型参数，而不是整个 Python 模型对象，通常更便于迁移与加载。 - 学习率调度器每个 Epoch 更新一次，与前面
T_max=epochs的定义一致。 - 这里根据验证准确率选择最佳模型；如果业务更关注类别均衡或误分类成本，可以改用 Macro-F1 或其他验证指标。
checkpoint_path = Path("best_resnet18.pth")
best_val_acc = -1.0
history = {
"train_loss": [],
"val_loss": [],
"val_accuracy": [],
"lr": [],
}
for epoch in range(epochs):
current_lr = optimizer.param_groups[0]["lr"]
train_loss = train_one_epoch(epoch, model, train_loader)
val_metrics = evaluate(model, val_loader, split_name="Validation")
history["train_loss"].append(train_loss)
history["val_loss"].append(val_metrics["loss"])
history["val_accuracy"].append(val_metrics["accuracy"])
history["lr"].append(current_lr)
print(
f"Epoch {epoch + 1}/{epochs} | "
f"Train Loss: {train_loss:.4f} | "
f"Val Loss: {val_metrics['loss']:.4f} | "
f"Val Accuracy: {val_metrics['accuracy']:.4%} | "
f"LR: {current_lr:.2e}"
)
# 只依据验证集选择模型，不使用测试集参与模型选择
if val_metrics["accuracy"] > best_val_acc:
best_val_acc = val_metrics["accuracy"]
torch.save(model.state_dict(), checkpoint_path)#覆盖写入 只保存一个
print(f"已保存当前最佳模型: {checkpoint_path.resolve()}")
scheduler.step()
print(f"训练完成，最佳验证准确率: {best_val_acc:.4%}")
13. 可视化训练曲线
训练损失持续下降通常说明模型正在拟合训练集；如果训练损失下降但验证损失上升，可能出现过拟合。曲线需要结合验证准确率与数据集规模综合判断。
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
axes[0].plot(range(1, epochs + 1), history["train_loss"], label="Train Loss")
axes[0].plot(range(1, epochs + 1), history["val_loss"], label="Validation Loss")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("Loss")
axes[0].set_title("Training and Validation Loss")
axes[0].legend()
axes[0].grid(True, alpha=0.3)#透明度0.3
axes[1].plot(
range(1, epochs + 1),
history["val_accuracy"],
label="Validation Accuracy"
)
axes[1].set_xlabel("Epoch")
axes[1].set_ylabel("Accuracy")
axes[1].set_title("Validation Accuracy")
axes[1].legend()
axes[1].grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
14. 加载最佳模型，并进行一次最终测试
测试集只在模型选择完成后使用一次，作为最终泛化性能评估。测试结果不应再用于反复调整超参数，否则测试集也会间接参与模型选择。
map_location=device：将保存的参数映射到当前运行设备。- 最终报告测试损失与准确率。
- 若要进一步分析错误，可基于返回的
targets和preds计算混淆矩阵、每类 Precision / Recall / F1；这些指标可以作为后续扩展。
if not checkpoint_path.is_file():
raise FileNotFoundError(
f"找不到最佳模型文件：{checkpoint_path.resolve()}。请先成功运行训练单元。"
)
model.load_state_dict(torch.load(checkpoint_path, map_location=device))
test_metrics = evaluate(model, test_loader, split_name="Test")
print(f"Test Loss: {test_metrics['loss']:.4f}")
print(f"Test Accuracy: {test_metrics['accuracy']:.4%}")
15. 预测单张图片
复用与验证集一致的预处理流程，对单张图片进行推理。将 image_path 改成实际存在的图片路径即可。
Image.open(...).convert("RGB")确保输入为三通道 RGB 图像。unsqueeze(0)添加 Batch 维度，使形状从[C, H, W]变为[1, C, H, W]。argmax取得最大 logit 对应的类别索引；若需要置信度，可额外对 logits 做 Softmax。
from PIL import Image
def predict_image(image_path, model):
image_path = Path(image_path)
if not image_path.is_file():
raise FileNotFoundError(f"找不到图片：{image_path.resolve()}")
image = Image.open(image_path).convert("RGB")
tensor = eval_tfs(image).unsqueeze(0).to(device)
model.eval()
with torch.inference_mode():
logits = model(tensor)
probabilities = torch.softmax(logits, dim=1)
confidence, pred_idx = probabilities.max(dim=1)# 获取每个样本的最高预测概率（置信度）及其对应的类别索引
return {
"class_name": categories[pred_idx.item()],
"confidence": confidence.item(),
}
# result = predict_image("./test/apple/example.jpg", model)
# print(result)
转载自 CSDN-专业IT技术社区
原文链接：https://blog.csdn.net/qq_41775119/article/details/168559756
