
2. 解冻部分或全部骨干层,并配合更小的学习率进行微调,使特征表示更贴合当前图像领域。
如果数据量较小,直接解冻全部参数很容易出现过拟合;如果当前图像领域与预训练数据差异很大,只训练分类头又可能不足以获得理想效果。因此,冻结范围、学习率设置和数据增强策略都应基于验证集表现来调整,而不应机械套用固定模板。
### 输入尺寸与归一化
EfficientNetB0 的预训练权重通常都默认输入图像沿用其原始训练时使用的尺寸设置与归一化方案。实际在 torchvision 中调用官方权重对象时,最稳妥的做法是直接使用其推荐的 transforms,而不是手动照搬均值、标准差或裁剪规则——这些细节一旦不一致,模型效果往往会受到明显影响。下面的代码示例采用的是当前 torchvision 常见接口;如果本地版本存在差异,应以已安装版本对应的官方文档和类型提示为准。
准备环境
建议先创建独立 Python 环境,再安装 PyTorch、torchvision 和 Pillow。由于 CPU、CUDA 以及其他硬件加速后端的安装方式可能不同,应根据目标机器在 PyTorch 官方安装页面选择对应构建版本。
```bash
python -m venv .venvsource .venv/bin/activate
python -m pip install --upgrade pippip install torch torchvision pillow
```Windows PowerShell 下的激活命令为:```powershell.venvScriptsActivate.ps1
```建议优先确认设备选择逻辑,而不是直接把 `cuda` 写死在代码里:```pythonimport torch
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")print("device:", device)```
整理数据集
`ImageFolder` 要求数据目录结构大致如下:
```text
dataset/train/
class_a/001.jpg
class_b/002.jpg
val/class_a/
101.jpgclass_b/
102.jpg```
验证集不能简单通过复制训练集图片得到。如果同一条视频、同一次连拍,或者同一张原图的多个裁剪版本同时出现在训练集和验证集中,那么验证结果通常会被高估。划分数据时应尽量按拍摄批次、个体来源或原始文件维度进行分组,具体边界则要结合实际采集方式判断。
如果当前只有一个总图片目录,建议先按来源进行分组,再执行训练集与验证集划分,而不是随机复制文件。划分完成后,还要检查每个类别的样本数量;当类别极度不均衡时,可以考虑使用加权损失、分层采样或补充少数类样本,不要只关注整体准确率这一项指标。
构建数据加载器
下面的实现方式会先从预训练权重中获取推荐的图像预处理流程,再在训练集上加入适度的随机增强。验证集则只保留确定性变换,以确保同一模型在不同轮次验证时接收到一致输入,便于稳定比较效果。
```python
from pathlib import Pathfrom torch.utils.data import DataLoader
from torchvision import datasets, transformsfrom torchvision.models import EfficientNet_B0_Weights
root = Path("dataset")
weights = EfficientNet_B0_Weights.DEFAULTbase_transform = weights.transforms()
train_transform = transforms.Compose([
transforms.RandomResizedCrop(base_transform.crop_size[0], scale=(0.75, 1.0)),transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),transforms.ToTensor(),
transforms.Normalize(base_transform.mean, base_transform.std),])
val_transform = base_transform
train_set = datasets.ImageFolder(root / "train", transform=train_transform)
val_set = datasets.ImageFolder(root / "val", transform=val_transform)if train_set.class_to_idx != val_set.class_to_idx:raise ValueError("训练集和验证集的类别映射不一致")
train_loader = DataLoader(train_set, batch_size=32, shuffle=True, num_workers=2)
val_loader = DataLoader(val_set, batch_size=32, shuffle=False, num_workers=2)print(train_set.classes)```
数据增强并不是越强越好。对于方向本身具有语义意义的图像分类任务,不应盲目加入水平翻转;对于颜色就是核心分类依据的场景,过强的颜色扰动也可能破坏标签含义。合理的数据增强应尽量模拟真实拍摄中的变化,而不是凭空制造数据集中本不存在的图像模式。
替换分类头并训练
EfficientNetB0 的分类器输出维度必须与当前数据集的类别数量保持一致。第一阶段通常先冻结特征提取层,只训练新的分类器:
```python
import torchfrom torch import nn
from torchvision.models import efficientnet_b0model = efficientnet_b0(weights=weights)for parameter in model.features.parameters():
parameter.requires_grad = Falsein_features = model.classifier[-1].in_featuresmodel.classifier[-1] = nn.Linear(in_features, len(train_set.classes))
model = model.to(device)criterion = nn.CrossEntropyLoss()optimizer = torch.optim.AdamW(
filter(lambda p: p.requires_grad, model.parameters()),lr=1e-3,
weight_decay=1e-4,)
```训练循环和验证循环应分别记录指标。验证阶段必须使用 `eval()` 与 `torch.no_grad()`,否则 BatchNorm、Dropout 或梯度计算都会影响最终评估结果。```pythondef run_epoch(model, loader, criterion, optimizer=None):
training = optimizer is not Nonemodel.train(training)
total_loss = 0.0total_correct = 0
total_count = 0context = torch.enable_grad() if training else torch.no_grad()with context:
for images, labels in loader:images, labels = images.to(device), labels.to(device)
logits = model(images)loss = criterion(logits, labels)
if training:
optimizer.zero_grad(set_to_none=True)loss.backward()
optimizer.step()total_loss = loss.item() * labels.size(0)total_correct = (logits.argmax(dim=1) == labels).sum().item()
total_count = labels.size(0)return total_loss / total_count, total_correct / total_countbest_val_acc = -1.0for epoch in range(10):
train_loss, train_acc = run_epoch(model, train_loader, criterion, optimizer)val_loss, val_acc = run_epoch(model, val_loader, criterion)
print(f"epoch={epoch 1} "
f"train_loss={train_loss:.4f} train_acc={train_acc:.4f} "f"val_loss={val_loss:.4f} val_acc={val_acc:.4f}"
)if val_acc > best_val_acc:best_val_acc = val_acc
torch.save({"model": model.state_dict(),
"classes": train_set.classes,}, "best.pt")
```这里的 `10` 只是示例中的训练轮数,并不意味着适用于所有图像分类数据集。实际训练时,应结合训练损失、验证损失以及类别级指标综合判断是否继续训练。若训练准确率持续上升,而验证损失同步上升,通常说明过拟合风险在增加,此时可以考虑减少解冻范围、增强正则化、补充数据,或引入早停策略。## 第二阶段微调当分类头训练趋于稳定后,可以进一步解冻最后一部分特征层,并将学习率降低一个数量级进行微调。具体解冻多少层,应根据数据量大小以及领域差异程度灵活调整:```pythonfor parameter in model.features[-2:].parameters():
parameter.requires_grad = Trueoptimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()),
lr=1e-4,weight_decay=1e-4,
)```
在开始微调前,应重新确认当前模型已经加载了验证集表现最好的检查点。若解冻后继续使用原来的高学习率,可能会破坏预训练得到的有效特征;而如果始终完全不解冻,又可能无法充分适应特殊背景、成像设备差异或类别间细粒度区别。
评估与导出
整体准确率并不能完整反映每个类别的识别表现。至少应保存混淆矩阵,用于分析哪些类别容易互相混淆;当数据类别分布不均衡时,还应重点查看每类的 precision、recall 和 F1 分数。若模型要用于实际业务筛选,还应明确不同错分类型的成本,不能仅凭准确率来选择最终模型。
模型推理阶段必须复用验证时的预处理流程以及类别映射:
```python
from PIL import Imagecheckpoint = torch.load("best.pt", map_location=device)model.load_state_dict(checkpoint["model"])
model.eval()image = val_transform(Image.open("sample.jpg").convert("RGB"))with torch.no_grad():
probabilities = model(image.unsqueeze(0).to(device)).softmax(dim=1)index = int(probabilities.argmax(dim=1).item())print({
"label": checkpoint["classes"][index],"confidence": float(probabilities[0, index]),
})```
分类概率并不等同于经过校准后的真实正确率。对于低置信度样本,可以将其送入人工复核流程,但阈值应通过独立验证数据来设定。如果后续需要部署到移动端或嵌入式设备,还应根据目标运行时选择 TorchScript、ONNX 或其他导出格式,并额外验证导出模型的输入输出一致性。
常见问题
### 验证准确率异常高
优先检查是否存在数据泄漏、重复图片、同一拍摄批次被划分到不同数据集,以及验证集预处理是否意外包含了随机增强。同时也要确认标签目录没有被错误复制或映射混乱。
### 损失出现 NaN
这通常可能与学习率过高、输入图片损坏、数值精度设置不当或标签异常有关。建议先用少量批次运行并检查输入张量的 `isfinite()`,再尝试降低学习率、临时关闭混合精度,逐步定位问题来源。
### 所有图片都预测为同一类
先检查类别目录结构和 `class_to_idx` 是否正确,确认不存在空目录、拼写变体或类别映射错误;再检查类别是否严重失衡、归一化是否与预训练权重匹配,以及训练过程中分类头参数是否确实得到了更新。
### 训练集准确率高但真实图片效果差
这通常说明训练数据分布与真实输入分布不一致,或者模型学习到了背景、拍摄设备、拍摄环境等非目标特征。应补充更多不同场景样本,按采集来源划分验证集,并通过可视化或遮挡实验检查模型真正关注的区域。
### 为什么不直接训练全部参数
当然可以直接训练全部参数,但这并不适合所有数据规模和图像分类任务。对于小数据集,端到端训练更容易过拟合,训练成本也更高。先冻结骨干再逐步微调,更有利于观察不同训练阶段的变化;而当数据量足够大、领域差异足够明显时,直接端到端训练也可能成为更合适的选择。
总结
一个真正具备落地能力的图像分类项目,关键从来不只是选用了 EfficientNetB0 这一模型,更重要的是将数据划分、图像预处理、迁移学习训练、模型评估与导出部署整套流程串联成可复现的闭环。进入实际应用阶段后,尤其要守住四条核心边界:训练集和验证集绝不能发生泄漏;训练阶段与推理阶段必须使用一致的变换;保存模型时必须同步保存类别映射;模型选择最终要以业务目标对应的评价指标为准。
在这些基础环节稳定之后,再去讨论更复杂的数据增强、类别重采样、模型量化或移动端部署,才会拥有明确且可信的比较基线。所有结论都应建立在独立验证数据和真实输入分布之上,而不是把单次训练日志当作普遍规律。
