简介:迁移学习是深度学习领域的一项关键技术,其核心原理是利用在大规模数据集(如ImageNet)上预训练好的模型权重,作为新任务的起点。这种方法通过复用模型已学到的通用视觉特征,避免了从零开始训练的巨大计算开销,显著提升了模型在小数据集上的收敛速度和最终性能,具有极高的技术价值。在计算机视觉的诸多应用场景中,如图像分类、目标检测等,迁移学习已成为标准实践方案。本文聚焦于一个具体的工程实践:使用经典的ResNet50预训练模型,结合PyTorch框架,对华为垃圾数据集进行细粒度图像分类。内容将详细拆解从虚拟环境配置、数据加载与预处理、模型微调、训练策略到最终模型评估与部署的全流程,为开发者提供一份可直接复用的实战指南,帮助读者快速掌握迁移学习在解决实际分类问题中的关键步骤与常见陷阱。
1. 项目缘起与核心价值
最近在整理一些计算机视觉的实战项目,发现很多朋友对“迁移学习”这个概念既熟悉又陌生。熟悉是因为这个词在各种教程里频繁出现,陌生是因为真要自己动手,从零开始用预训练模型去解决一个具体的分类问题,中间还是有不少坑要踩。正好,我之前用华为开源的垃圾数据集,基于经典的ResNet50模型,完整跑通了一个图像分类系统。这个项目麻雀虽小,五脏俱全,涵盖了从环境搭建、数据处理、模型微调、训练技巧到最终部署测试的全流程。它特别适合那些已经学过一些Python和深度学习基础,但还没亲手把一个模型从“能用”做到“好用”的开发者。今天,我就把这个项目的核心思路、关键代码以及我趟过的那些“坑”详细拆解一遍,你可以把它看作一份可以直接“抄作业”的实战指南。
为什么选这个组合?ResNet50作为CNN领域的常青树,其残差结构有效缓解了深度网络的梯度消失问题,在ImageNet上预训练的权重是一个极强的视觉特征提取器,为我们提供了一个高起点。而华为的垃圾数据集,包含了可回收物、厨余垃圾、有害垃圾等类别,是一个典型的、有实际意义的细粒度图像分类任务。通过迁移学习,我们无需从随机初始化权重开始训练,极大地节省了计算资源和时间,尤其适合在个人电脑或算力有限的场景下,快速构建一个性能不错的分类器。这个项目的源码,就是教你如何把这两者高效、稳定地结合起来。
2. 环境搭建与依赖管理:避开第一个大坑
动手之前,环境是地基。很多教程会轻描淡写地让你pip install一堆包,但版本冲突往往是项目跑不起来的第一元凶。我的经验是,为每个项目创建独立的虚拟环境是必须养成的习惯。
2.1 创建并激活虚拟环境
我强烈推荐使用conda进行环境管理,它能更好地处理一些复杂的C++依赖(比如PyTorch的CUDA版本)。如果你没有安装Anaconda或Miniconda,先去官网下载安装。之后,打开终端(Windows用Anaconda Prompt,Linux/macOS用终端),执行以下命令:
# 创建一个名为‘garbage_classification’的Python3.8环境 conda create -n garbage_classification python=3.8 # 激活这个环境 conda activate garbage_classification选择Python 3.8是因为它在深度学习生态中兼容性非常广泛,是一个比较稳妥的版本。环境激活后,你的命令行提示符前面应该会出现(garbage_classification)的字样。
2.2 安装核心依赖库
接下来安装核心的深度学习框架和工具。这里以PyTorch为例,因为它和Torchvision的配合非常紧密,对于计算机视觉任务特别友好。你需要根据自己是否有NVIDIA显卡以及CUDA版本来选择安装命令。可以去PyTorch官网生成对应的命令。假设我们使用CUDA 11.3,命令如下:
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113如果你没有GPU,或者想先确保环境能跑通,可以安装CPU版本:
pip install torch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1安装完PyTorch后,继续安装其他必要的库:
pip install numpy pandas matplotlib opencv-python pillow scikit-learn tqdm tensorboard这里解释一下每个库的作用:numpy和pandas用于数据处理;matplotlib用于可视化;opencv-python和Pillow是图像处理的左膀右臂;scikit-learn用于计算评估指标;tqdm可以给你的循环加上美观的进度条;tensorboard则是训练过程可视化的利器。
注意:安装时可能会遇到网络问题导致某些包下载慢或失败。一个实用的技巧是使用国内镜像源,例如在pip命令后加上
-i https://pypi.tuna.tsinghua.edu.cn/simple。但安装PyTorch时,务必使用官方--extra-index-url或从官网获取命令,因为镜像源可能不包含带CUDA版本的PyTorch。
2.3 验证环境与常见问题
安装完成后,写一个简单的脚本来验证环境是否正常:
import torch import torchvision print(f“PyTorch版本: {torch.__version__}“) print(f“CUDA是否可用: {torch.cuda.is_available()}“) if torch.cuda.is_available(): print(f“当前CUDA设备: {torch.cuda.get_device_name(0)}“)如果CUDA可用,会显示你的显卡型号。如果不可用但你有显卡,那大概率是PyTorch版本和CUDA驱动版本不匹配,需要去NVIDIA控制面板查看CUDA版本,然后重新安装对应版本的PyTorch。这是环境配置中最常见的一个坑。
3. 数据准备与预处理:模型效果的基石
模型再强大,如果喂给它的数据是“脏”的或者组织混乱的,效果也会大打折扣。华为垃圾数据集通常以文件夹结构组织,每个类别的图片放在一个以类别名命名的子文件夹里。我们的任务就是把它转换成模型训练需要的格式。
3.1 数据集目录结构解析与加载
假设你的数据集解压后结构如下:
huawei_garbage_dataset/ ├── train/ │ ├── cardboard/ # 纸板类图片 │ ├── glass/ # 玻璃类图片 │ ├── metal/ # 金属类图片 │ └── ... # 其他类别 └── val/ # 验证集,结构同train ├── cardboard/ ├── glass/ └── ...PyTorch提供了torchvision.datasets.ImageFolder这个非常方便的工具,它能自动根据这种目录结构加载数据,并为每个子文件夹分配一个标签。首先,我们需要定义数据变换(Transforms)。这是预处理的核心,目的是将原始图片转换成张量(Tensor),并进行标准化,使数据分布更利于模型收敛。
from torchvision import transforms # 定义训练集的数据增强和变换 train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪并缩放到224x224 transforms.RandomHorizontalFlip(), # 随机水平翻转,简单有效的增强 transforms.RandomRotation(10), # 随机旋转±10度 transforms.ColorJitter(brightness=0.2, contrast=0.2), # 随机调整亮度对比度 transforms.ToTensor(), # 将PIL图像或numpy数组转换为张量 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 标准化 ]) # 定义验证集/测试集的变换(通常不进行数据增强) val_transform = transforms.Compose([ transforms.Resize(256), # 将短边缩放到256 transforms.CenterCrop(224), # 从中心裁剪224x224 transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])这里有几个关键点需要解释:第一,输入尺寸为什么是224x224?因为ResNet50以及大多数在ImageNet上预训练的模型,其全连接层输入特征维度是基于224x224的图像设计的,这是一个标准尺寸。第二,标准化使用的均值[0.485, 0.456, 0.406]和标准差[0.229, 0.224, 0.225]是ImageNet数据集上百万张图片统计出来的全局均值与标准差。使用预训练模型时,必须采用相同的标准化参数,因为模型的权重是在这个数据分布下学习到的。第三,训练集使用了多种数据增强(RandomHorizontalFlip, RandomRotation, ColorJitter),这是为了防止模型过拟合到训练集,提升其泛化能力。而验证集不需要增强,我们希望看到模型在原始数据上的真实表现。
3.2 创建数据加载器(DataLoader)
使用ImageFolder加载数据,并用DataLoader包装,它负责在训练时按批次(batch)提供数据,并支持多进程数据加载以加速IO。
from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader # 路径替换成你自己的数据集路径 train_dataset = ImageFolder(root=‘./huawei_garbage_dataset/train‘, transform=train_transform) val_dataset = ImageFolder(root=‘./huawei_garbage_dataset/val‘, transform=val_transform) # 创建数据加载器 batch_size = 32 # 根据你的GPU内存调整,太小训练慢,太大可能爆内存 num_workers = 4 # 用于数据加载的子进程数,可以加快数据读取速度 train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True) # 打印一些基本信息 print(f“训练集类别数: {len(train_dataset.classes)}“) print(f“训练集样本数: {len(train_dataset)}“) print(f“验证集样本数: {len(val_dataset)}“) print(f“类别名称: {train_dataset.classes}“)shuffle=True意味着每个epoch(完整遍历一次训练集)开始时,训练数据的顺序会被打乱,这有助于模型学习更通用的特征,避免依赖于数据顺序。pin_memory=True是一个优化选项,当你的数据从CPU转移到GPU时,如果数据在锁页内存中,转移速度会更快,这在有GPU的情况下建议开启。
3.3 数据可视化与检查
在开始训练前,花几分钟看一眼你的数据是非常有必要的。这能帮你发现一些潜在问题,比如图片损坏、标签错误、或者数据增强的效果是否符合预期。
import matplotlib.pyplot as plt import numpy as np # 获取一个批次的数据 images, labels = next(iter(train_loader)) # 将张量转换回图片格式显示 def imshow(inp, title=None): “”“从张量显示图像。”“” inp = inp.numpy().transpose((1, 2, 0)) # 从(C, H, W)转换为(H, W, C) mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) inp = std * inp + mean # 反标准化 inp = np.clip(inp, 0, 1) # 将像素值限制在[0,1]之间 plt.imshow(inp) if title is not None: plt.title(title) plt.pause(0.001) # 暂停一下让绘图更新 # 显示一个批次中的部分图片 fig = plt.figure(figsize=(12, 8)) for i in range(min(8, batch_size)): # 显示前8张 ax = fig.add_subplot(2, 4, i+1) ax.axis(‘off‘) ax.set_title(train_dataset.classes[labels[i]]) imshow(images[i]) plt.show()这个步骤能直观地确认数据加载和预处理流程是否正确。如果你看到图片扭曲得不成样子,或者类别标签明显不对,就需要回头检查数据集的目录结构或变换流程了。
4. ResNet50模型加载与迁移学习改造
这是项目的核心环节。我们不是从零训练ResNet50,而是利用它在ImageNet上学习到的强大特征提取能力,只训练最后几层,使其适应我们的垃圾分类任务。
4.1 加载预训练模型并冻结底层参数
首先,我们加载在ImageNet上预训练好的ResNet50模型。torchvision.models提供了非常便捷的接口。
import torch.nn as nn import torchvision.models as models # 加载预训练的ResNet50模型,并下载预训练权重 model = models.resnet50(pretrained=True) # 冻结模型的所有参数,在初始阶段不让它们参与梯度更新 for param in model.parameters(): param.requires_grad = Falsepretrained=True会自动下载并加载预训练权重。接着,我们通过循环将模型中所有参数的requires_grad属性设置为False。这意味着在反向传播时,这些参数不会计算梯度,也就不会被优化器更新。这样做是因为模型的前面几层(卷积层)学习到的是非常通用的低级特征(如边缘、纹理、颜色),这些特征对于我们的垃圾图像识别任务同样有用,我们不需要改变它们。
4.2 替换最后的全连接层
ResNet50原模型的最后一层是一个1000个神经元的全连接层(对应ImageNet的1000个类别)。我们的垃圾数据集可能只有几个或几十个类别,所以必须替换这一层。
# 获取原全连接层的输入特征数 num_ftrs = model.fc.in_features # 假设我们的垃圾数据集有6个类别(例如:纸板、玻璃、金属、纸张、塑料、其他) num_classes = len(train_dataset.classes) # 这里用之前从数据集中获取的类别数 # 用一个新的全连接层替换原来的fc层 # 这个新的层默认 requires_grad=True model.fc = nn.Linear(num_ftrs, num_classes) # 将新替换的fc层,以及我们后面可能想要微调的层,设置为可训练 for param in model.fc.parameters(): param.requires_grad = True这里,model.fc就是ResNet50最后的全连接层。我们新建了一个nn.Linear层,输入维度保持不变(num_ftrs,对于ResNet50是2048),输出维度改为我们的类别数num_classes。只有新加入的这层参数是需要从头开始训练的。
4.3 选择性地微调中间层(进阶技巧)
对于某些与预训练任务差异较大的数据集,或者当我们有相对充足的数据时,可以尝试解冻模型靠后的部分卷积层,让它们也进行微调(Fine-tuning),以学习更贴合新任务的特征。一个常见的策略是解冻最后两个“阶段”(stage)的层。
# 以ResNet50为例,其结构分为多个阶段(layer1, layer2, layer3, layer4) # 我们解冻layer4和layer3的参数 for name, param in model.named_parameters(): if “layer4“ in name or “layer3“ in name: param.requires_grad = True这样做的好处是,模型可以调整更深层的、更抽象的特征表示来适应新任务,可能获得比只训练最后一层更好的性能。但风险是如果新数据量很小,很容易导致过拟合。我的建议是:先冻结所有层,只训练最后的全连接层,得到一个基准模型。如果验证集性能达到瓶颈,再尝试解冻后面几层进行微调,并配合更小的学习率和更强的正则化(如Dropout)。
5. 训练策略、损失函数与优化器配置
模型准备好了,数据也加载好了,接下来就是定义如何训练它。这包括选择损失函数来衡量预测与真实标签的差距,选择优化器来更新模型参数,以及制定学习率调整策略。
5.1 损失函数与优化器选择
对于多分类任务,交叉熵损失(Cross-Entropy Loss)是标准选择。优化器方面,Adam因其自适应学习率特性,在大多数情况下都能取得不错的效果且收敛快,非常适合作为默认选择。
import torch.optim as optim # 将模型移动到GPU(如果可用) device = torch.device(“cuda:0“ if torch.cuda.is_available() else “cpu“) model = model.to(device) # 定义损失函数 criterion = nn.CrossEntropyLoss() # 定义优化器,只优化那些 requires_grad=True 的参数 # 初始学习率设置为一个较小的值,因为我们是微调,不是从头训练 optimizer = optim.Adam(model.parameters(), lr=0.001) # 也可以使用SGD,有时配合动量(momentum)能获得更好的最终精度,但需要仔细调整学习率 # optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)这里有一个关键细节:model.parameters()传递给优化器后,优化器只会更新那些requires_grad=True的参数。这正是我们想要的——只更新我们新加的fc层(以及可能解冻的层)。
5.2 学习率调度器(Learning Rate Scheduler)
在训练过程中动态调整学习率非常重要。一开始可以用较大的学习率快速下降,后期则需要调小学习率以便精细调整,收敛到更好的局部最优点。ReduceLROnPlateau是一个很实用的策略:当验证集指标(如loss)在连续几个epoch内不再下降时,自动降低学习率。
from torch.optim import lr_scheduler # 当验证损失连续3个epoch没有下降时,将学习率乘以0.1 scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode=‘min‘, factor=0.1, patience=3, verbose=True)mode=‘min‘表示我们监控的指标(这里是验证损失)越低越好。patience=3给了模型3个epoch的“耐心期”,如果连续3个epoch验证损失都不降,就触发学习率衰减。verbose=True会在控制台打印学习率变化的信息。
5.3 训练循环的完整实现
训练循环是深度学习的引擎,它反复执行前向传播、计算损失、反向传播、参数更新这个过程。一个健壮且功能完整的训练循环还应包含验证、日志记录和模型保存。
import time import copy from tqdm import tqdm # 用于显示进度条 def train_model(model, criterion, optimizer, scheduler, num_epochs=25): since = time.time() best_model_wts = copy.deepcopy(model.state_dict()) # 保存最优模型的权重 best_acc = 0.0 history = {‘train_loss‘: [], ‘train_acc‘: [], ‘val_loss‘: [], ‘val_acc‘: []} for epoch in range(num_epochs): print(f‘Epoch {epoch}/{num_epochs - 1}‘) print(‘-‘ * 10) # 每个epoch都有训练和验证阶段 for phase in [‘train‘, ‘val‘]: if phase == ‘train‘: model.train() # 设置模型为训练模式(启用Dropout, BatchNorm更新) dataloader = train_loader else: model.eval() # 设置模型为评估模式(关闭Dropout, 固定BatchNorm统计量) dataloader = val_loader running_loss = 0.0 running_corrects = 0 # 使用tqdm包装数据加载器,显示进度条 for inputs, labels in tqdm(dataloader, desc=phase): inputs = inputs.to(device) labels = labels.to(device) # 梯度清零 optimizer.zero_grad() # 前向传播 # 只在训练阶段追踪计算图以计算梯度 with torch.set_grad_enabled(phase == ‘train‘): outputs = model(inputs) _, preds = torch.max(outputs, 1) # 获取预测类别(最大值的索引) loss = criterion(outputs, labels) # 只在训练阶段进行反向传播和优化 if phase == ‘train‘: loss.backward() optimizer.step() # 统计 running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / len(dataloader.dataset) epoch_acc = running_corrects.double() / len(dataloader.dataset) # 记录历史数据 if phase == ‘train‘: history[‘train_loss‘].append(epoch_loss) history[‘train_acc‘].append(epoch_acc.item()) # 学习率调度器在验证阶段根据验证损失更新,所以这里不调用 else: history[‘val_loss‘].append(epoch_loss) history[‘val_acc‘].append(epoch_acc.item()) scheduler.step(epoch_loss) # 根据验证损失调整学习率 print(f‘{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}‘) # 深度拷贝模型权重(如果验证准确率更高) if phase == ‘val‘ and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = copy.deepcopy(model.state_dict()) # 可以在这里保存当前最好的模型 torch.save(model.state_dict(), f‘best_model_epoch_{epoch}.pth‘) print(f‘==> 保存新的最佳模型,准确率: {best_acc:.4f}‘) print() time_elapsed = time.time() - since print(f‘训练完成,用时 {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s‘) print(f‘最佳验证准确率: {best_acc:.4f}‘) # 加载最佳模型权重 model.load_state_dict(best_model_wts) return model, history这个训练函数包含了几个关键实践:第一,使用model.train()和model.eval()来切换模型模式,这对BatchNorm和Dropout层的行为至关重要。第二,使用torch.set_grad_enabled()上下文管理器来控制是否计算梯度,在验证阶段禁用可以节省内存。第三,在验证阶段结束后,根据验证损失调用scheduler.step()。第四,持续追踪并保存验证集上性能最好的模型权重,这是一种简单的模型选择策略,能确保我们最终得到的是泛化能力最强的模型,而不是最后一个epoch可能过拟合的模型。
6. 模型评估、可视化与错误分析
训练完成后,我们不能只看最后的准确率数字就完事。需要从多个维度评估模型,理解它在哪里做得好,在哪里容易出错,这能为我们后续的改进提供方向。
6.1 在测试集上评估最终模型
首先,我们需要一个独立的测试集(如果数据集提供了的话),或者从验证集中留出一部分作为测试集,来最终评估模型的泛化能力。评估过程与验证阶段类似。
def evaluate_model(model, test_loader): model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for inputs, labels in tqdm(test_loader, desc=‘Evaluating‘): inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) _, preds = torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) # 计算整体准确率 from sklearn.metrics import accuracy_score, classification_report, confusion_matrix accuracy = accuracy_score(all_labels, all_preds) print(f‘测试集整体准确率: {accuracy:.4f}‘) # 打印详细的分类报告(精确率、召回率、F1分数) print(“\n分类报告:“) print(classification_report(all_labels, all_preds, target_names=test_dataset.classes)) # 计算混淆矩阵 cm = confusion_matrix(all_labels, all_preds) return cm, accuracy, all_preds, all_labels分类报告能清晰地展示每个类别的精确率(Precision)、召回率(Recall)和F1分数(F1-Score)。如果某个类别的分数明显偏低,说明模型在这个类别上识别困难,可能需要针对性增加该类的数据或进行数据增强。
6.2 可视化训练过程与混淆矩阵
可视化能帮助我们直观理解模型的训练动态和错误模式。
import matplotlib.pyplot as plt import seaborn as sns from sklearn.metrics import confusion_matrix def plot_training_history(history): fig, axes = plt.subplots(1, 2, figsize=(12, 4)) # 绘制损失曲线 axes[0].plot(history[‘train_loss‘], label=‘Train Loss‘) axes[0].plot(history[‘val_loss‘], label=‘Val Loss‘) axes[0].set_title(‘Training and Validation Loss‘) axes[0].set_xlabel(‘Epoch‘) axes[0].set_ylabel(‘Loss‘) axes[0].legend() axes[0].grid(True) # 绘制准确率曲线 axes[1].plot(history[‘train_acc‘], label=‘Train Acc‘) axes[1].plot(history[‘val_acc‘], label=‘Val Acc‘) axes[1].set_title(‘Training and Validation Accuracy‘) axes[1].set_xlabel(‘Epoch‘) axes[1].set_ylabel(‘Accuracy‘) axes[1].legend() axes[1].grid(True) plt.tight_layout() plt.show() def plot_confusion_matrix(cm, class_names): plt.figure(figsize=(10, 8)) sns.heatmap(cm, annot=True, fmt=‘d‘, cmap=‘Blues‘, xticklabels=class_names, yticklabels=class_names) plt.title(‘Confusion Matrix‘) plt.ylabel(‘True Label‘) plt.xlabel(‘Predicted Label‘) plt.tight_layout() plt.show() # 使用保存的历史数据绘图 plot_training_history(history) # 使用评估函数返回的混淆矩阵绘图 plot_confusion_matrix(cm, test_dataset.classes)通过损失和准确率曲线,我们可以判断训练是否正常。理想情况下,训练损失和验证损失都应该稳步下降并最终趋于平稳,训练准确率和验证准确率同步上升。如果出现训练损失持续下降但验证损失上升(即“过拟合”),或者两者都很高且不下降(“欠拟合”),就需要调整模型或数据。混淆矩阵则能具体显示模型把哪些类别互相混淆了,比如“塑料瓶”是否容易被误判为“玻璃瓶”,这为数据收集和模型改进提供了明确目标。
6.3 可视化模型预测结果(Grad-CAM)
对于图像分类模型,我们常常想知道模型到底是根据图像的哪一部分做出决策的。Grad-CAM(梯度加权类激活映射)是一种可视化技术,可以生成一个热力图,高亮显示对模型预测贡献最大的图像区域。
import cv2 import numpy as np class GradCAM: def __init__(self, model, target_layer): self.model = model self.target_layer = target_layer self.gradients = None self.activations = None self._register_hooks() def _register_hooks(self): def forward_hook(module, input, output): self.activations = output def backward_hook(module, grad_input, grad_output): self.gradients = grad_output[0] self.target_layer.register_forward_hook(forward_hook) self.target_layer.register_full_backward_hook(backward_hook) def generate_cam(self, input_image, target_class=None): self.model.eval() output = self.model(input_image) if target_class is None: target_class = output.argmax(dim=1).item() self.model.zero_grad() output[0, target_class].backward() # 计算权重 weights = self.gradients.mean(dim=(2, 3), keepdim=True) # 全局平均池化梯度 cam = (weights * self.activations).sum(dim=1, keepdim=True) cam = torch.relu(cam) # ReLU操作,只保留对类别有正向贡献的特征 cam = cam - cam.min() cam = cam / cam.max() return cam.squeeze().cpu().detach().numpy() # 使用示例:对ResNet50,我们通常取最后一个卷积层(layer4) target_layer = model.layer4[-1].conv3 # ResNet50最后一个bottleneck的最后一个卷积层 grad_cam = GradCAM(model, target_layer) # 获取一张测试图片 img, label = test_dataset[0] input_tensor = img.unsqueeze(0).to(device) # 增加batch维度 # 生成CAM cam = grad_cam.generate_cam(input_tensor) cam = cv2.resize(cam, (224, 224)) # 调整到输入图像大小 # 将CAM叠加到原图上 img_np = img.numpy().transpose(1, 2, 0) img_np = np.clip(img_np * np.array([0.229, 0.224, 0.225]) + np.array([0.485, 0.456, 0.406]), 0, 1) # 反标准化 heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET) heatmap = np.float32(heatmap) / 255 superimposed_img = heatmap * 0.4 + np.float32(img_np) * 0.6 plt.figure(figsize=(10, 5)) plt.subplot(1, 3, 1) plt.imshow(img_np) plt.title(f‘Original (True: {test_dataset.classes[label]})‘) plt.axis(‘off‘) plt.subplot(1, 3, 2) plt.imshow(cam, cmap=‘jet‘) plt.title(‘Grad-CAM Heatmap‘) plt.axis(‘off‘) plt.subplot(1, 3, 3) plt.imshow(superimposed_img) plt.title(‘Overlay‘) plt.axis(‘off‘) plt.show()通过Grad-CAM可视化,你可以判断模型是否关注了正确的物体区域。例如,在垃圾分类中,如果模型识别“易拉罐”时,热图集中在罐体的金属部分和拉环上,那是合理的;如果热图集中在背景上,那说明模型可能学到了错误的特征关联,需要警惕。
7. 模型部署与简易推理脚本
训练好的模型最终要用来做预测。我们需要编写一个推理脚本,能够加载保存的模型权重,并对新的单张图片或一批图片进行分类。
7.1 保存与加载模型
训练时我们保存了最佳模型的权重(.pth文件)。推理时需要重新实例化模型结构,然后加载权重。
def load_trained_model(model_path, num_classes): “”“加载训练好的模型”“” # 1. 实例化模型结构(必须和训练时完全一致) model = models.resnet50(pretrained=False) # 不加载ImageNet权重 num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, num_classes) # 2. 加载训练好的权重 model.load_state_dict(torch.load(model_path, map_location=device)) model = model.to(device) model.eval() # 设置为评估模式 return model # 假设我们有6个类别 num_classes = 6 trained_model = load_trained_model(‘best_model_epoch_10.pth‘, num_classes)这里有一个非常重要的细节:torch.load时指定了map_location=device。这能确保模型权重被加载到正确的设备上(CPU或GPU),避免因为训练和推理环境设备不同而报错。
7.2 单张图片推理函数
这个函数接收一张图片的路径,完成从读取、预处理到预测的全流程。
from PIL import Image def predict_single_image(image_path, model, class_names, transform): “”“对单张图片进行预测”“” # 1. 读取和预处理图片 image = Image.open(image_path).convert(‘RGB‘) # 确保是三通道 image_tensor = transform(image).unsqueeze(0) # 应用变换并增加batch维度 image_tensor = image_tensor.to(device) # 2. 预测 with torch.no_grad(): outputs = model(image_tensor) probabilities = torch.nn.functional.softmax(outputs, dim=1) # 转换为概率 confidence, predicted_idx = torch.max(probabilities, 1) # 3. 返回结果 predicted_class = class_names[predicted_idx.item()] confidence_score = confidence.item() return predicted_class, confidence_score # 使用示例 class_names = [‘cardboard‘, ‘glass‘, ‘metal‘, ‘paper‘, ‘plastic‘, ‘trash‘] # 替换为你的类别名 val_transform = transforms.Compose([...]) # 使用和验证集相同的变换 img_path = ‘./test_image.jpg‘ pred_class, confidence = predict_single_image(img_path, trained_model, class_names, val_transform) print(f‘预测类别: {pred_class}, 置信度: {confidence:.2%}‘)7.3 批量推理与结果导出
在实际应用中,我们可能需要对一个文件夹下的所有图片进行批量分类。
import os from pathlib import Path def predict_batch_images(image_dir, model, class_names, transform, output_csv=‘predictions.csv‘): “”“对一个目录下的所有图片进行批量预测并保存结果”“” results = [] image_extensions = [‘.jpg‘, ‘.jpeg‘, ‘.png‘, ‘.bmp‘] image_paths = [p for p in Path(image_dir).iterdir() if p.suffix.lower() in image_extensions] for img_path in tqdm(image_paths, desc=‘Predicting‘): try: pred_class, confidence = predict_single_image(str(img_path), model, class_names, transform) results.append({ ‘image_path‘: str(img_path), ‘predicted_class‘: pred_class, ‘confidence‘: confidence }) except Exception as e: print(f“处理图片 {img_path} 时出错: {e}“) results.append({ ‘image_path‘: str(img_path), ‘predicted_class‘: ‘ERROR‘, ‘confidence‘: 0.0 }) # 保存为CSV文件 import pandas as pd df = pd.DataFrame(results) df.to_csv(output_csv, index=False, encoding=‘utf-8-sig‘) # 使用utf-8-sig避免中文乱码 print(f“预测完成,结果已保存至 {output_csv}“) return df # 使用示例 predictions_df = predict_batch_images(‘./test_images/‘, trained_model, class_names, val_transform)这个批量推理脚本还包含了简单的错误处理,防止某张损坏的图片导致整个程序中断。输出为CSV格式,方便后续用Excel或其他工具进行分析。
8. 项目总结与进阶优化思路
走完以上所有步骤,一个基于ResNet50和迁移学习的垃圾图像分类系统就完整搭建起来了。从环境配置、数据处理、模型构建、训练调优到评估部署,我们覆盖了一个深度学习项目的主要生命周期。这个过程里,有几个点是我在实际操作中体会特别深的。
第一,数据质量决定上限。模型和算法只是逼近这个上限的工具。在华为垃圾数据集中,如果某些类别的图片数量严重不足,或者图片质量参差不齐(模糊、遮挡、光照差异大),模型的性能天花板就会被拉低。因此,在数据准备阶段花时间做清洗、做增强,甚至自己收集补充一些数据,其投资回报率往往比后期拼命调参要高得多。例如,可以尝试更复杂的数据增强,如MixUp、CutMix,或者使用AutoAugment等自动增强策略。
第二,学习率是超参数之王。在微调预训练模型时,学习率的设置尤为关键。对于新添加的全连接层,我们可以给一个相对较大的学习率(如0.01),而对于预训练模型中解冻的层,学习率应该设置得更小(如0.001或0.0001),通常称为“差分学习率”。这可以通过优化器的参数组来实现:
# 为模型的不同部分设置不同的学习率 optimizer = optim.SGD([ {‘params‘: model.layer4.parameters(), ‘lr‘: 0.001}, # 解冻层,小学习率微调 {‘params‘: model.fc.parameters(), ‘lr‘: 0.01} # 新层,较大学习率 ], momentum=0.9)第三,不要忽视正则化的力量。当解冻更多层进行微调时,模型容量变大,过拟合风险增加。除了数据增强,在模型中引入Dropout层或者在优化器中加入权重衰减(Weight Decay)都是有效的正则化手段。例如,可以在全连接层后加入Dropout:model.fc = nn.Sequential(nn.Dropout(0.5), nn.Linear(num_ftrs, num_classes))。
第四,模型集成是提升性能的利器。如果计算资源允许,训练多个不同初始化或不同超参数的模型,然后将它们的预测结果进行平均(对于概率)或投票(对于类别),通常能获得比单一模型更稳定、更准确的结果。这对于竞赛或对精度要求极高的场景尤其有效。
这个项目代码本身是一个完整的、可运行的系统,但它更是一个模板和起点。你可以轻松地将其适配到其他的图像分类任务上,比如花卉分类、皮肤病识别、工业品缺陷检测等,只需要更换数据集和调整最后的类别数。深度学习的魅力就在于这种强大的可迁移性,而掌握了一套完整的项目流程,你就拥有了快速解决新问题的能力。
本文还有配套的精品资源,点击获取