公司动态

基于ResNet50与迁移学习的图像分类实战:从环境搭建到模型部署

📅 2026/8/28 8:59:34
基于ResNet50与迁移学习的图像分类实战:从环境搭建到模型部署
简介迁移学习是深度学习领域的一项关键技术其核心原理是利用在大规模数据集如ImageNet上预训练好的模型权重作为新任务的起点。这种方法通过复用模型已学到的通用视觉特征避免了从零开始训练的巨大计算开销显著提升了模型在小数据集上的收敛速度和最终性能具有极高的技术价值。在计算机视觉的诸多应用场景中如图像分类、目标检测等迁移学习已成为标准实践方案。本文聚焦于一个具体的工程实践使用经典的ResNet50预训练模型结合PyTorch框架对华为垃圾数据集进行细粒度图像分类。内容将详细拆解从虚拟环境配置、数据加载与预处理、模型微调、训练策略到最终模型评估与部署的全流程为开发者提供一份可直接复用的实战指南帮助读者快速掌握迁移学习在解决实际分类问题中的关键步骤与常见陷阱。1. 项目缘起与核心价值最近在整理一些计算机视觉的实战项目发现很多朋友对“迁移学习”这个概念既熟悉又陌生。熟悉是因为这个词在各种教程里频繁出现陌生是因为真要自己动手从零开始用预训练模型去解决一个具体的分类问题中间还是有不少坑要踩。正好我之前用华为开源的垃圾数据集基于经典的ResNet50模型完整跑通了一个图像分类系统。这个项目麻雀虽小五脏俱全涵盖了从环境搭建、数据处理、模型微调、训练技巧到最终部署测试的全流程。它特别适合那些已经学过一些Python和深度学习基础但还没亲手把一个模型从“能用”做到“好用”的开发者。今天我就把这个项目的核心思路、关键代码以及我趟过的那些“坑”详细拆解一遍你可以把它看作一份可以直接“抄作业”的实战指南。为什么选这个组合ResNet50作为CNN领域的常青树其残差结构有效缓解了深度网络的梯度消失问题在ImageNet上预训练的权重是一个极强的视觉特征提取器为我们提供了一个高起点。而华为的垃圾数据集包含了可回收物、厨余垃圾、有害垃圾等类别是一个典型的、有实际意义的细粒度图像分类任务。通过迁移学习我们无需从随机初始化权重开始训练极大地节省了计算资源和时间尤其适合在个人电脑或算力有限的场景下快速构建一个性能不错的分类器。这个项目的源码就是教你如何把这两者高效、稳定地结合起来。2. 环境搭建与依赖管理避开第一个大坑动手之前环境是地基。很多教程会轻描淡写地让你pip install一堆包但版本冲突往往是项目跑不起来的第一元凶。我的经验是为每个项目创建独立的虚拟环境是必须养成的习惯。2.1 创建并激活虚拟环境我强烈推荐使用conda进行环境管理它能更好地处理一些复杂的C依赖比如PyTorch的CUDA版本。如果你没有安装Anaconda或Miniconda先去官网下载安装。之后打开终端Windows用Anaconda PromptLinux/macOS用终端执行以下命令# 创建一个名为‘garbage_classification’的Python3.8环境 conda create -n garbage_classification python3.8 # 激活这个环境 conda activate garbage_classification选择Python 3.8是因为它在深度学习生态中兼容性非常广泛是一个比较稳妥的版本。环境激活后你的命令行提示符前面应该会出现(garbage_classification)的字样。2.2 安装核心依赖库接下来安装核心的深度学习框架和工具。这里以PyTorch为例因为它和Torchvision的配合非常紧密对于计算机视觉任务特别友好。你需要根据自己是否有NVIDIA显卡以及CUDA版本来选择安装命令。可以去PyTorch官网生成对应的命令。假设我们使用CUDA 11.3命令如下pip install torch1.12.1cu113 torchvision0.13.1cu113 torchaudio0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113如果你没有GPU或者想先确保环境能跑通可以安装CPU版本pip install torch1.12.1 torchvision0.13.1 torchaudio0.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(brightness0.2, contrast0.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‘, transformtrain_transform) val_dataset ImageFolder(root‘./huawei_garbage_dataset/val‘, transformval_transform) # 创建数据加载器 batch_size 32 # 根据你的GPU内存调整太小训练慢太大可能爆内存 num_workers 4 # 用于数据加载的子进程数可以加快数据读取速度 train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workersnum_workers, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_sizebatch_size, shuffleFalse, num_workersnum_workers, pin_memoryTrue) # 打印一些基本信息 print(f“训练集类别数: {len(train_dataset.classes)}“) print(f“训练集样本数: {len(train_dataset)}“) print(f“验证集样本数: {len(val_dataset)}“) print(f“类别名称: {train_dataset.classes}“)shuffleTrue意味着每个epoch完整遍历一次训练集开始时训练数据的顺序会被打乱这有助于模型学习更通用的特征避免依赖于数据顺序。pin_memoryTrue是一个优化选项当你的数据从CPU转移到GPU时如果数据在锁页内存中转移速度会更快这在有GPU的情况下建议开启。3.3 数据可视化与检查在开始训练前花几分钟看一眼你的数据是非常有必要的。这能帮你发现一些潜在问题比如图片损坏、标签错误、或者数据增强的效果是否符合预期。import matplotlib.pyplot as plt import numpy as np # 获取一个批次的数据 images, labels next(iter(train_loader)) # 将张量转换回图片格式显示 def imshow(inp, titleNone): “”“从张量显示图像。”“” 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, i1) 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(pretrainedTrue) # 冻结模型的所有参数在初始阶段不让它们参与梯度更新 for param in model.parameters(): param.requires_grad FalsepretrainedTrue会自动下载并加载预训练权重。接着我们通过循环将模型中所有参数的requires_grad属性设置为False。这意味着在反向传播时这些参数不会计算梯度也就不会被优化器更新。这样做是因为模型的前面几层卷积层学习到的是非常通用的低级特征如边缘、纹理、颜色这些特征对于我们的垃圾图像识别任务同样有用我们不需要改变它们。4.2 替换最后的全连接层ResNet50原模型的最后一层是一个1000个神经元的全连接层对应ImageNet的1000个类别。我们的垃圾数据集可能只有几个或几十个类别所以必须替换这一层。# 获取原全连接层的输入特征数 num_ftrs model.fc.in_features # 假设我们的垃圾数据集有6个类别例如纸板、玻璃、金属、纸张、塑料、其他 num_classes len(train_dataset.classes) # 这里用之前从数据集中获取的类别数 # 用一个新的全连接层替换原来的fc层 # 这个新的层默认 requires_gradTrue 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_gradTrue 的参数 # 初始学习率设置为一个较小的值因为我们是微调不是从头训练 optimizer optim.Adam(model.parameters(), lr0.001) # 也可以使用SGD有时配合动量momentum能获得更好的最终精度但需要仔细调整学习率 # optimizer optim.SGD(model.parameters(), lr0.001, momentum0.9)这里有一个关键细节model.parameters()传递给优化器后优化器只会更新那些requires_gradTrue的参数。这正是我们想要的——只更新我们新加的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‘, factor0.1, patience3, verboseTrue)mode‘min‘表示我们监控的指标这里是验证损失越低越好。patience3给了模型3个epoch的“耐心期”如果连续3个epoch验证损失都不降就触发学习率衰减。verboseTrue会在控制台打印学习率变化的信息。5.3 训练循环的完整实现训练循环是深度学习的引擎它反复执行前向传播、计算损失、反向传播、参数更新这个过程。一个健壮且功能完整的训练循环还应包含验证、日志记录和模型保存。import time import copy from tqdm import tqdm # 用于显示进度条 def train_model(model, criterion, optimizer, scheduler, num_epochs25): 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, descphase): 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_namestest_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, annotTrue, fmt‘d‘, cmap‘Blues‘, xticklabelsclass_names, yticklabelsclass_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_classNone): self.model.eval() output self.model(input_image) if target_class is None: target_class output.argmax(dim1).item() self.model.zero_grad() output[0, target_class].backward() # 计算权重 weights self.gradients.mean(dim(2, 3), keepdimTrue) # 全局平均池化梯度 cam (weights * self.activations).sum(dim1, keepdimTrue) 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(pretrainedFalse) # 不加载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_locationdevice)) 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_locationdevice。这能确保模型权重被加载到正确的设备上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, dim1) # 转换为概率 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, indexFalse, 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} # 新层较大学习率 ], momentum0.9)第三不要忽视正则化的力量。当解冻更多层进行微调时模型容量变大过拟合风险增加。除了数据增强在模型中引入Dropout层或者在优化器中加入权重衰减Weight Decay都是有效的正则化手段。例如可以在全连接层后加入Dropoutmodel.fc nn.Sequential(nn.Dropout(0.5), nn.Linear(num_ftrs, num_classes))。第四模型集成是提升性能的利器。如果计算资源允许训练多个不同初始化或不同超参数的模型然后将它们的预测结果进行平均对于概率或投票对于类别通常能获得比单一模型更稳定、更准确的结果。这对于竞赛或对精度要求极高的场景尤其有效。这个项目代码本身是一个完整的、可运行的系统但它更是一个模板和起点。你可以轻松地将其适配到其他的图像分类任务上比如花卉分类、皮肤病识别、工业品缺陷检测等只需要更换数据集和调整最后的类别数。深度学习的魅力就在于这种强大的可迁移性而掌握了一套完整的项目流程你就拥有了快速解决新问题的能力。本文还有配套的精品资源点击获取