芯耀
与非AI
353
（注：使用文中代码的时候需要加上冒号后的缩进，否则报语法错误。如需复制代码建议点击文末左下角阅读原文。）
1.项目背景
在当前全球倡导绿色低碳生活的大环境下，生活垃圾的精准分类与高效回收已成为智慧城市建设中不可或缺的一环。然而，面对海量且种类繁杂的废弃物，单纯依赖人工分拣不仅效率低下，且极易受主观判断偏差的影响，导致回收链路的后端处理成本激增。
为了攻克这一痛点，本项目立足于计算机视觉技术，深度挖掘来自 Kaggle 的大规模垃圾分类数据集。该数据集构建了一个极其丰富的真实场景库，通过 15,000 张高清晰度图像，全方位覆盖了从日常塑料制品、金属罐体到纺织品、有机废弃物等 30 个细分维度，为模型提供了极佳的特征学习空间。我们希望通过训练一个高性能的深度学习分类器，探索如何利用轻量化神经网络在保证识别精度的前提下，实现对各类可回收材料与家居废弃物的快速感官模拟，从而为自动分拣设备和智能回收终端提供核心的技术支撑，真正让“变废为宝”走向自动化与智能化。
2.数据集介绍
本实验数据集来源于Kaggle，该数据集包含15,000张图像（每张256x256像素），涵盖30个不同类别的各种可回收材料、一般垃圾和家居用品。每个类别包含500张图像，每个子类别包含250张图像，为垃圾分类和回收领域的研究与开发提供了丰富多样的资源。通过提供大量高质量图像，该数据集旨在支持构建稳健、准确的垃圾分拣和分类系统。
该数据集涵盖多种废弃物类别和物品，包括：
塑料：此类别包括塑料水瓶、汽水瓶、洗涤剂瓶、购物袋、垃圾袋、食品容器、一次性餐具、吸管和杯盖的图片。这些物品占家庭塑料垃圾的很大一部分，对回收利用至关重要。
纸张和纸板：此类别包含报纸、办公用纸、杂志、纸箱和纸板包装的图片。这些物品通常都可以回收利用，在减少森林砍伐和保护自然资源方面发挥着至关重要的作用。
玻璃：此类别包含玻璃材质的饮料瓶、食品罐和化妆品容器的图片。玻璃是一种高度可回收的材料，正确的分类和分拣对于有效的回收流程至关重要。
金属：此类别包含铝制汽水罐、铝制食品罐、钢制食品罐和气雾罐的图片。金属废料具有很高的回收价值，如果能够正确识别和分类，就可以高效地进行处理。
有机废弃物：此类别包含食物废弃物的图片，例如果皮、蔬菜残渣、蛋壳、咖啡渣和茶包。有机废弃物可以堆肥或用于沼气生产，从而减轻垃圾填埋场的负担并产生宝贵的资源。
纺织品：此类别包含服装和鞋类的图片。纺织品废弃物日益增多，令人担忧，正确的分类有助于回收利用，并减少时尚产业对环境的影响。
3.技术工具
Python版本:3.9
代码编辑器：jupyter notebook
4.实验过程
4.1导入数据
在开始深度学习项目之前，首先需要准备好环境。这一步我们导入了项目所需的全部库，包括 PyTorch 核心库、Torchvision图像处理库、Scikit-learn 评估工具以及数据可视化相关的库。
import torchvision.transforms.functional
from PIL import Image
from torch.utils.data import Dataset
from torchvision.transforms import v2
from torch.utils.data import DataLoader
import torch
import os
import random
import matplotlib.pyplot as plt
import numpy as np
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import mobilenet_v3_small # 使用轻量级模型 MobileNetV3
from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report, ConfusionMatrixDisplay
运行
自定义数据集加载类
为了适配本项目中“垃圾分类”数据集的存储结构（包含 default 和 real_world 子文件夹），我们继承了 Dataset 类自定义了 WasteDataset。该类负责扫描文件路径、按照 6:2:2 的比例划分训练集、验证集和测试集，并实现数据的实时读取与预处理。
# 准备自定义数据集类
classWasteDataset(Dataset):
def__init__(self, root_dir, split, transform=None):
"""
root_dir: 数据集根目录
split: 指定数据集类型 ('train', 'val', 'test')
transform: 图像预处理/增强操作
"""
self.root_dir = root_dir
self.transform = transform
self.classes = sorted(os.listdir(root_dir)) # 获取类别名称并排序
self.image_paths = [] # 存储图像绝对路径
self.labels = [] # 存储对应的标签索引
# 遍历每个类别文件夹
for i, class_name inenumerate(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 == 'test'
image_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):
# 返回数据集的总样本数
returnlen(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(brightness=0.5, contrast=0.5, saturation=0.2), # 颜色抖动
v2.RandomAffine(degrees=5, translate=(0.1, 0.1), scale=(0.8, 1.3),
interpolation=torchvision.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), # 转换数据类型为 float32
v2.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, shuffle=True, num_workers=int(os.cpu_count()*0.8))
val_data_loader = DataLoader(val_dataset, batch_size, shuffle=False, num_workers=int(os.cpu_count()*0.2))
test_data_loader = DataLoader(test_dataset, batch_size, shuffle=True)
# 统一封装加载器
data_loaders = {
"train": train_data_loader,
"val": val_data_loader,
"test": test_data_loader
}
运行
4.2数据可视化
为了直观地查看加载后的图像数据及其对应的标签，我们编写了一个可视化函数 visualize_batch。由于在预处理阶段我们对图像进行了标准化（Normalization），在显示图像前，需要进行反标准化处理，将像素值还原到 [0, 1]或 [0, 255] 的范围。
defvisualize_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 inrange(0, batch_size):
# 创建 4x4 的子图布局
ax = plt.subplot(4, 4, i + 1)
# 将张量从 (C, H, W) 转换为 (H, W, C) 以适配 matplotlib
image = batch["image"][i].cpu().numpy()
image = image.transpose((1, 2, 0))
# 执行反标准化操作：image = (normalized_image * std) + mean
image = 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 个类别）。
classWasteClassificationModel(nn.Module):
def__init__(self):
super().__init__()
# 加载预训练的 MobileNetV3-Small 模型权重
self.mobnet = mobilenet_v3_small(weights=torchvision.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'}
)
# 自定义卷积层：进一步提取特征，输入通道数为 576（MobileNet 该层输出），输出为 300
self.conv1 = nn.Conv2d(576, 300, 3)
# 全连接层：将提取的特征映射到 30 个类别空间
# 10800 是经过卷积和展平（flatten）后的特征维度
self.fc1 = nn.Linear(10800, 30)
# Dropout 层：在训练过程中随机失活神经元，防止过拟合
self.dr = nn.Dropout()
defforward(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_dim=1)
# 4. 应用 Dropout
x = self.dr(x)
# 5. 最后通过全连接层输出分类结果（Logits）
output = self.fc1(x)
return output
运行
要点说明：
为什么选择 MobileNetV3？ 它的参数量小、计算速度快，非常适合移动端或嵌入式设备的垃圾分类场景。
特征提取器的作用： 我们不从零开始训练，而是“站在巨人的肩膀上”，利用已经学会识别形状、纹理的预训练层，只训练最后几层来适应我们的特定数据集。
4.4训练模型
在定义好模型结构后，我们需要配置训练环境。这包括选择计算设备、定义损失函数（CrossEntropyLoss）和优化器（Adam）。为了得到性能最强的模型，我们会在训练过程中实时监控训练集和验证集的准确率与损失，并保存最优的模型权重。
# 检查 GPU 是否可用，优先使用 GPU 加速训练
device = torch.device("cuda:0"if torch.cuda.is_available() else"cpu")
print(f"Using Device {device}")
# --- 准备损失函数与优化器 ---
model = WasteClassificationModel() # 实例化模型
model = model.to(device) # 将模型移动到计算设备
criterian = nn.CrossEntropyLoss() # 多分类任务常用的交叉熵损失函数
model_optimizer = optim.Adam(model.parameters(), lr=0.001) # Adam 优化器
num_epochs = 50# 迭代轮数
# 初始化最优指标，用于保存表现最好的模型
best_acc = {"train": -1000, "val": -1000}
best_loss = {"train": 1000, "val": 1000}
# 定义模型权重保存路径
best_accuracy_model_path = {
"train": os.path.join("/kaggle/working/", "train_acc_best.pt"),
"val": os.path.join("/kaggle/working/", "val_acc_best.pt")
}
best_loss_model_path = {
"train": os.path.join("/kaggle/working/", "train_loss_best.pt"),
"val": os.path.join("/kaggle/working/", "val_loss_best.pt")
}
# --- 开始训练循环 ---
for epoch inrange(num_epochs):
print(f'Epoch {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 inenumerate(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_acc
torch.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_loss
torch.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(phase=="train")：这是一个非常优雅的处理方式，它确保了在验证阶段不计算梯度，从而节省显存并加快运行速度。
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 inenumerate(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_names=class_names))
# 5. 计算并绘制混淆矩阵
# 混淆矩阵可以直观展示模型将哪些类别误分类为了其他类别
cm = confusion_matrix(y_true, y_preds)
cm_disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=class_names)
# 创建画布，并根据类别数量调整大小
fig, ax = plt.subplots(figsize=(16, 14))
cm_disp.plot(ax=ax) # 绘制混淆矩阵图
plt.xticks(rotation=45) # 旋转横坐标标签，防止重叠
plt.show()
运行
4.6模型预测
为了验证模型在完全未见的数据上的表现，我们从测试集（Test Set）中抽取一个批次的图像进行推理。通过将原始图像还原显示，并在标题中标注真实类别与预测类别，我们可以清晰地观察到模型在实际应用中的分类准确性。
# 从测试加载器中获取一个批次的测试数据
test_batch = next(iter(data_loaders["test"]))
# 将图像和标签移动到计算设备（GPU 或 CPU）
inputs, 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 inrange(len(predictions)):
# 创建 8行2列 的子图布局（对应 batch_size 为 16 的情况）
ax = plt.subplot(8, 2, i + 1)
# 图像后处理：从 Tensor 转回 Numpy
image = 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 = f"True: {true_label}nPred: {predicted_label}"
# 如果预测错误，可以用不同颜色或格式标注（此处统一显示）
plt.title(title_label, fontsize=10)
plt.axis("off") # 隐藏坐标轴
# 优化布局，防止标题重叠
plt.tight_layout()
plt.show()
运行
5.总结
本实验基于 Kaggle 提供的包含 15,000 张高质量图像的垃圾分类数据集，成功构建并验证了一个涵盖塑料、纸张、玻璃、金属、有机废弃物及纺织品等 30 个类别的深度学习分类系统。实验结果显示，模型在多类别任务上表现稳健，最终达到了 81% 的整体准确率（Accuracy）。通过分类报告可以看出，模型在特征鲜明的类别上表现卓越，如咖啡渣（Coffee Grounds）和食物残渣（Food Waste）的 F1 分数均突破了 0.9，展现了极强的识别精度；同时，在塑料洗涤剂瓶和一次性餐具等工业化形态明显的物品上也保持了极高的精确率。
尽管模型在部分相似材质（如纸板包装与纸箱）的区分上仍存在一定的提升空间，但整体评估指标充分证明了利用迁移学习结合MobileNetV3架构能够有效处理复杂的现实废弃物分类任务，为构建自动化、智能化的绿色回收分拣系统提供了可靠的技术支持与实验参考。
