深度学习实战-基于迁移学习MobileNet的可回收物与生活垃圾图像分类模型
1.项目背景在当前全球倡导绿色低碳生活的大环境下生活垃圾的精准分类与高效回收已成为智慧城市建设中不可或缺的一环。然而面对海量且种类繁杂的废弃物单纯依赖人工分拣不仅效率低下且极易受主观判断偏差的影响导致回收链路的后端处理成本激增。为了攻克这一痛点本项目立足于计算机视觉技术深度挖掘来自 Kaggle 的大规模垃圾分类数据集。该数据集构建了一个极其丰富的真实场景库通过 15,000 张高清晰度图像全方位覆盖了从日常塑料制品、金属罐体到纺织品、有机废弃物等 30 个细分维度为模型提供了极佳的特征学习空间。我们希望通过训练一个高性能的深度学习分类器探索如何利用轻量化神经网络在保证识别精度的前提下实现对各类可回收材料与家居废弃物的快速感官模拟从而为自动分拣设备和智能回收终端提供核心的技术支撑真正让“变废为宝”走向自动化与智能化。2.数据集介绍本实验数据集来源于Kaggle该数据集包含15,000张图像每张256x256像素涵盖30个不同类别的各种可回收材料、一般垃圾和家居用品。每个类别包含500张图像每个子类别包含250张图像为垃圾分类和回收领域的研究与开发提供了丰富多样的资源。通过提供大量高质量图像该数据集旨在支持构建稳健、准确的垃圾分拣和分类系统。该数据集涵盖多种废弃物类别和物品包括塑料此类别包括塑料水瓶、汽水瓶、洗涤剂瓶、购物袋、垃圾袋、食品容器、一次性餐具、吸管和杯盖的图片。这些物品占家庭塑料垃圾的很大一部分对回收利用至关重要。纸张和纸板此类别包含报纸、办公用纸、杂志、纸箱和纸板包装的图片。这些物品通常都可以回收利用在减少森林砍伐和保护自然资源方面发挥着至关重要的作用。玻璃此类别包含玻璃材质的饮料瓶、食品罐和化妆品容器的图片。玻璃是一种高度可回收的材料正确的分类和分拣对于有效的回收流程至关重要。金属此类别包含铝制汽水罐、铝制食品罐、钢制食品罐和气雾罐的图片。金属废料具有很高的回收价值如果能够正确识别和分类就可以高效地进行处理。有机废弃物此类别包含食物废弃物的图片例如果皮、蔬菜残渣、蛋壳、咖啡渣和茶包。有机废弃物可以堆肥或用于沼气生产从而减轻垃圾填埋场的负担并产生宝贵的资源。纺织品此类别包含服装和鞋类的图片。纺织品废弃物日益增多令人担忧正确的分类有助于回收利用并减少时尚产业对环境的影响。3.技术工具Python版本:3.9代码编辑器jupyter notebook4.实验过程4.1导入数据在开始深度学习项目之前首先需要准备好环境。这一步我们导入了项目所需的全部库包括 PyTorch 核心库、Torchvision 图像处理库、Scikit-learn 评估工具以及数据可视化相关的库。import torchvision.transforms.functionalfrom PIL import Imagefrom torch.utils.data import Datasetfrom torchvision.transforms import v2from torch.utils.data import DataLoaderimport torchimport osimport randomimport matplotlib.pyplot as pltimport numpy as npimport torch.optim as optimimport torch.nn as nnimport torch.nn.functional as Ffrom torchvision.models import mobilenet_v3_small # 使用轻量级模型 MobileNetV3from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_namesfrom sklearn.metrics import accuracy_score, confusion_matrix, classification_report, ConfusionMatrixDisplay为方便大家学习 这里给大家整理了一份学习资料包 需要的同学 根据下图自取即可自定义数据集加载类为了适配本项目中“垃圾分类”数据集的存储结构包含 default 和 real_world 子文件夹我们继承了 Dataset 类自定义了 WasteDataset。该类负责扫描文件路径、按照 6:2:2 的比例划分训练集、验证集和测试集并实现数据的实时读取与预处理。# 准备自定义数据集类class WasteDataset(Dataset):def __init__(self, root_dir, split, transformNone):root_dir: 数据集根目录split: 指定数据集类型 (train, val, test)transform: 图像预处理/增强操作self.root_dir root_dirself.transform transformself.classes sorted(os.listdir(root_dir)) # 获取类别名称并排序self.image_paths [] # 存储图像绝对路径self.labels [] # 存储对应的标签索引# 遍历每个类别文件夹for i, class_name in enumerate(self.classes):class_dir os.path.join(root_dir, class_name)# 遍历类别下的子文件夹默认图像和真实世界图像for subfolder in [default, real_world]:subfolder_dir os.path.join(class_dir, subfolder)image_names os.listdir(subfolder_dir)random.shuffle(image_names) # 随机打乱数据保证分配均匀# 按照比例划分数据集60% 训练, 20% 验证, 20% 测试if split train:image_names image_names[:int(0.6 * len(image_names))]elif split val:image_names image_names[int(0.6 * len(image_names)):int(0.8 * len(image_names))]else: # split testimage_names image_names[int(0.8 * len(image_names)):]# 将筛选后的路径及标签存入列表for image_name in image_names:self.image_paths.append(os.path.join(subfolder_dir, image_name))self.labels.append(i)def __len__(self):# 返回数据集的总样本数return len(self.image_paths)def __getitem__(self, index):# 获取单个样本image_path self.image_paths[index]label self.labels[index]# 使用 PIL 读取图像并转换为 RGB 模式image Image.open(image_path).convert(RGB)# 如果定义了预处理流程则应用if self.transform:image self.transform(image)# 以字典形式返回数据和标签data {image: image,label: label}return data定义数据增强与加载器在深度学习中数据增强可以有效防止模型过拟合。我们为训练集设计了包含亮度变化、随机仿射变换和高斯模糊的增强流程对于验证集和测试集则仅进行缩放和标准化处理以保证评估的准确性。# 准备数据集及数据加载器 (Dataloader)# 训练集数据增强包含亮度对比度抖动、随机旋转缩放、高斯模糊等train_pil_transform v2.Compose([v2.ColorJitter(brightness0.5, contrast0.5, saturation0.2), # 颜色抖动v2.RandomAffine(degrees5, translate(0.1, 0.1), scale(0.8, 1.3),interpolationtorchvision.transforms.InterpolationMode.BILINEAR), # 仿射变换v2.Resize(size(256, 256)), # 统一尺寸v2.GaussianBlur(kernel_size(7, 13), sigma(0.1, 0.2)), # 高斯模糊v2.PILToTensor(), # 转换为张量v2.ToDtype(torch.float32), # 转换数据类型为 float32v2.Normalize(mean(0.5, 0.5, 0.5), std(0.5, 0.5, 0.5)) # 标准化])# 验证集数据处理仅缩放与标准化val_pil_transform v2.Compose([v2.Resize(size(256, 256)),v2.PILToTensor(),v2.ToDtype(torch.float32),v2.Normalize(mean(0.5, 0.5, 0.5), std(0.5, 0.5, 0.5))])# 测试集数据处理保持与验证集一致test_pil_transform v2.Compose([v2.Resize(size(256, 256)),v2.PILToTensor(),v2.ToDtype(torch.float32),v2.Normalize(mean(0.5, 0.5, 0.5), std(0.5, 0.5, 0.5))])# 将预处理流程封装进字典data_transforms {train: train_pil_transform,val: val_pil_transform,test: test_pil_transform,}# 实例化 Dataset 类train_dataset WasteDataset(/kaggle/input/recyclable-and-household-waste-classification/images/images, train,data_transforms[train])val_dataset WasteDataset(/kaggle/input/recyclable-and-household-waste-classification/images/images, val,data_transforms[val])test_dataset WasteDataset(/kaggle/input/recyclable-and-household-waste-classification/images/images, test,data_transforms[test])# 统计数据集信息image_datasets {train: train_dataset,val: val_dataset,test: test_dataset}class_names train_dataset.classes # 类别名称dataset_sizes {x: len(image_datasets[x]) for x in [train, val]} # 数据集大小batch_size 16 # 设置批大小# 封装为 DataLoader利用多进程加速数据读取train_data_loader DataLoader(train_dataset, batch_size, shuffleTrue, num_workersint(os.cpu_count()*0.8))val_data_loader DataLoader(val_dataset, batch_size, shuffleFalse, num_workersint(os.cpu_count()*0.2))test_data_loader DataLoader(test_dataset, batch_size, shuffleTrue)# 统一封装加载器data_loaders {train: train_data_loader,val: val_data_loader,test: test_data_loader}4.2数据可视化为了直观地查看加载后的图像数据及其对应的标签我们编写了一个可视化函数 visualize_batch。由于在预处理阶段我们对图像进行了标准化Normalization在显示图像前需要进行反标准化处理将像素值还原到 [0, 1]或 [0, 255] 的范围。def visualize_batch(batch, classes, dataset_type):可视化一个批次的数据batch: DataLoader返回的一个批次数据classes: 类别名称列表dataset_type: 数据集类型名称如 train# 初始化画布设置标题和大小fig plt.figure({} batch.format(dataset_type),figsize(batch_size, batch_size))# 预处理时使用的均值和标准差用于反标准化mean np.array([0.5, 0.5, 0.5])std np.array([0.5, 0.5, 0.5])# 遍历当前批次中的所有图片for i in range(0, batch_size):# 创建 4x4 的子图布局ax plt.subplot(4, 4, i 1)# 将张量从 (C, H, W) 转换为 (H, W, C) 以适配 matplotlibimage batch[image][i].cpu().numpy()image image.transpose((1, 2, 0))# 执行反标准化操作image (normalized_image * std) meanimage std * image mean# 将像素值缩放回 [0, 255] 并转为 uint8 类型image (image * 255).astype(uint8)# 获取当前图像的标签索引并查表得到类别名称idx batch[label][i]label classes[idx]# 绘制图像设置标题并隐藏坐标轴plt.imshow(image)plt.title(label)plt.axis(off)# 自动调整子图间距并展示plt.tight_layout()plt.show()# ---------------------------------------------------------# 可视化训练集数据# ---------------------------------------------------------# 获取训练集中的第一个批次train_batch next(iter(data_loaders[train]))# 调用函数进行展示visualize_batch(train_batch, class_names, train)为方便大家学习 这里给大家整理了一份学习资料包 需要的同学 根据下图自取即可4.3构建模型在本项目中我们采用了迁移学习的策略。我们选择轻量级且高效的MobileNetV3-Small作为主干网络Backbone并利用其在 ImageNet 数据集上预训练好的权重来提取图像特征。通过 create_feature_extractor 灵活地截取模型中间层随后添加自定义的卷积层和全连接层以适配本项目的垃圾分类任务共 30 个类别。class WasteClassificationModel(nn.Module):def __init__(self):super().__init__()# 加载预训练的 MobileNetV3-Small 模型权重self.mobnet mobilenet_v3_small(weightstorchvision.models.MobileNet_V3_Small_Weights.IMAGENET1K_V1)# 获取模型节点名称可选用于查看网络结构确认层名称train_nodes, eval_nodes get_graph_node_names(self.mobnet)# 提取指定的中间层特征这里选择 features.12 层作为特征输出节点self.feature_extraction create_feature_extractor(self.mobnet, return_nodes{features.12: mob_feature})# 自定义卷积层进一步提取特征输入通道数为 576MobileNet 该层输出输出为 300self.conv1 nn.Conv2d(576, 300, 3)# 全连接层将提取的特征映射到 30 个类别空间# 10800 是经过卷积和展平flatten后的特征维度self.fc1 nn.Linear(10800, 30)# Dropout 层在训练过程中随机失活神经元防止过拟合self.dr nn.Dropout()def forward(self, x):前向传播过程# 1. 通过主干网络提取中间层特征feature_layer self.feature_extraction(x)[mob_feature]# 2. 经过自定义卷积层并使用 ReLU 激活函数x F.relu(self.conv1(feature_layer))# 3. 展平操作将多维特征图拉直为一维向量跳过 batch 维度x x.flatten(start_dim1)# 4. 应用 Dropoutx self.dr(x)# 5. 最后通过全连接层输出分类结果Logitsoutput self.fc1(x)return output要点说明为什么选择 MobileNetV3它的参数量小、计算速度快非常适合移动端或嵌入式设备的垃圾分类场景。特征提取器的作用我们不从零开始训练而是“站在巨人的肩膀上”利用已经学会识别形状、纹理的预训练层只训练最后几层来适应我们的特定数据集。4.4训练模型在定义好模型结构后我们需要配置训练环境。这包括选择计算设备GPU 或 CPU、定义损失函数CrossEntropyLoss和优化器Adam。为了得到性能最强的模型我们会在训练过程中实时监控训练集和验证集的准确率与损失并保存最优的模型权重。# --- 开始训练循环 ---for epoch in range(num_epochs):print(fEpoch {epoch}/{num_epochs - 1})print(- * 10)# 每个 epoch 包含 训练(train) 和 验证(val) 两个阶段for phase in [train, val]:if phase train:model.train() # 设置为训练模式启用 Dropout 等else:model.eval() # 设置为评估模式running_loss 0.0 # 统计累计损失running_corrects 0.0 # 统计预测正确的样本数# 遍历数据加载器中的批次数据for idx, data in enumerate(data_loaders[phase]):# 将数据和标签发送到计算设备inputs, labels data[image].to(device), data[label].to(device)# 梯度清零防止上一个 batch 的梯度干扰model_optimizer.zero_grad()# 只有在训练阶段才启用梯度计算with torch.set_grad_enabled(phase train):outputs model(inputs) # 前向传播_, preds torch.max(outputs, 1) # 获取预测概率最大的索引即类别loss criterian(outputs, labels) # 计算损失# 训练阶段执行反向传播和参数更新if phase train:loss.backward() # 反向传播计算梯度model_optimizer.step() # 更新模型参数# 累加统计数据running_loss loss.item() * inputs.size(0)running_corrects torch.sum(preds labels.data)# 计算当前 epoch 的平均损失和准确率epoch_loss running_loss / dataset_sizes[phase]epoch_acc running_corrects / dataset_sizes[phase]print(f{phase} Loss : {epoch_loss:.4f} Acc: {epoch_acc:.4f})# --- 模型保存逻辑 ---# 如果当前准确率高于历史最高值保存模型if epoch_acc best_acc[phase]:best_acc[phase] epoch_acctorch.save({epoch: epoch,model_state_dict: model.state_dict(),optimizer_state_dict: model_optimizer.state_dict(),loss: running_loss}, best_accuracy_model_path[phase])# 如果当前损失低于历史最低值保存模型if epoch_loss best_loss[phase]:best_loss[phase] epoch_losstorch.save({epoch: epoch,model_state_dict: model.state_dict(),optimizer_state_dict: model_optimizer.state_dict(),loss: running_loss}, best_loss_model_path[phase])要点说明torch.set_grad_enabled(phasetrain)这是一个非常优雅的处理方式它确保了在验证阶段不计算梯度从而节省显存并加快运行速度。Checkpoint机制代码中保存了 model_state_dict 和 optimizer_state_dict这允许我们在训练中断后重新加载并继续训练而不仅仅是保存了模型结果。4.5模型评估训练完成后我们需要在验证集或测试集上评估模型的泛化能力。这里我们加载了训练过程中损失值最低Best Loss的模型权重并利用 scikit-learn 库生成详细的分类指标报告以及可视化混淆矩阵。# 1. 初始化模型结构并加载训练好的最优权重model WasteClassificationModel()# 加载保存的 checkpoint 文件checkpoint torch.load(/kaggle/working/train_loss_best.pt)model.to(device) # 移动到 GPU/CPU# 从字典中提取模型状态字典并加载model.load_state_dict(checkpoint[model_state_dict])# 2. 设置为评估模式关闭 Dropout 等model.eval()# 准备容器记录预测值和真实标签用于后续评估data_size len(image_datasets[val])y_preds []y_true []# 3. 遍历验证集数据进行推理for idx, data in enumerate(data_loaders[val]):# 获取输入图像和标签inputs, labels data[image].to(device), data[label].to(device)# 禁用梯度计算以节省内存with torch.no_grad():outputs model(inputs)# 寻找每行最大概率值的索引即为预测类别_, predictions torch.max(outputs, 1)# 将结果转回 CPU 并存储到列表中y_preds.extend(predictions.cpu().numpy().tolist())y_true.extend(labels.cpu().numpy().tolist())# 4. 打印详细的分类报告 (包含精确率 Precision、召回率 Recall、F1 值)print(\n-------Classification Report-------\n)# target_names 参数允许我们将数字标签映射回原始的类别名称print(classification_report(y_true, y_preds, target_namesclass_names))# 5. 计算并绘制混淆矩阵# 混淆矩阵可以直观展示模型将哪些类别误分类为了其他类别cm confusion_matrix(y_true, y_preds)cm_disp ConfusionMatrixDisplay(confusion_matrixcm, display_labelsclass_names)# 创建画布并根据类别数量调整大小fig, ax plt.subplots(figsize(16, 14))cm_disp.plot(axax) # 绘制混淆矩阵图plt.xticks(rotation45) # 旋转横坐标标签防止重叠plt.show()为方便大家学习 这里给大家整理了一份学习资料包 需要的同学 根据下图自取即可4.6模型预测为了验证模型在完全未见的数据上的表现我们从测试集Test Set中抽取一个批次的图像进行推理。通过将原始图像还原显示并在标题中标注真实类别与预测类别我们可以清晰地观察到模型在实际应用中的分类准确性。# 从测试加载器中获取一个批次的测试数据test_batch next(iter(data_loaders[test]))# 将图像和标签移动到计算设备GPU 或 CPUinputs, labels test_batch[image].to(device), test_batch[label].to(device)# 模型推理关闭梯度计算以提高效率with torch.no_grad():outputs model(inputs)# 获取预测概率最大的索引即预测的类别编号_, predictions torch.max(outputs.data, 1)# --- 结果可视化展示 ---# 创建画布设置适当的大小以便观察fig plt.figure(Test Batch Results, figsize(batch_size, batch_size))mean np.array([0.5, 0.5, 0.5]) # 反标准化用的均值std np.array([0.5, 0.5, 0.5]) # 反标准化用的标准差# 遍历当前批次中的所有预测结果for i in range(len(predictions)):# 创建 8行2列 的子图布局对应 batch_size 为 16 的情况ax plt.subplot(8, 2, i 1)# 图像后处理从 Tensor 转回 Numpyimage inputs[i].cpu().numpy()# 将通道维度从第一位移到最后一位 (C, H, W) - (H, W, C)image image.transpose((1, 2, 0))# 反标准化还原像素值image std * image mean# 缩放至 [0, 255] 并转换为 uint8 格式以便 plt 显示image (image * 255).astype(uint8)# 获取该图像对应的真实标签名称idx labels[i]true_label class_names[idx]# 获取模型预测的标签名称predicted_label class_names[predictions[i]]# 绘制图像plt.imshow(image)# 在标题中对比真实值和预测值title_label fTrue: {true_label}\nPred: {predicted_label}# 如果预测错误可以用不同颜色或格式标注此处统一显示plt.title(title_label, fontsize10)plt.axis(off) # 隐藏坐标轴# 优化布局防止标题重叠plt.tight_layout()plt.show()5.总结本实验基于 Kaggle 提供的包含 15,000 张高质量图像的垃圾分类数据集成功构建并验证了一个涵盖塑料、纸张、玻璃、金属、有机废弃物及纺织品等 30 个类别的深度学习分类系统。实验结果显示模型在多类别任务上表现稳健最终达到了 81% 的整体准确率Accuracy。通过分类报告可以看出模型在特征鲜明的类别上表现卓越如咖啡渣Coffee Grounds和食物残渣Food Waste的 F1 分数均突破了 0.9展现了极强的识别精度同时在塑料洗涤剂瓶和一次性餐具等工业化形态明显的物品上也保持了极高的精确率。尽管模型在部分相似材质如纸板包装与纸箱的区分上仍存在一定的提升空间但整体评估指标充分证明了利用迁移学习结合 MobileNetV3 架构能够有效处理复杂的现实废弃物分类任务为构建自动化、智能化的绿色回收分拣系统提供了可靠的技术支持与实验参考。为方便大家学习 这里给大家整理了一份学习资料包 需要的同学 根据下图自取即可