基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统
今天在查资料的时候在网上看到一篇文章，博主是基于TensorFlow实现的CNN来完成对垃圾分类识别的，想到最近正好在使用Pytorch就想也做一下，就当是项目开发实践了。先看下动态操作效果：
原文在这里，截图如下所示：
感兴趣的话可以自行移步原文。
原文实现的垃圾分类系统更为详尽，一共细分了40个不同的子类，如下所示：
{
"0": "其他垃圾/一次性快餐盒",
"1": "其他垃圾/污损塑料",
"2": "其他垃圾/烟蒂",
"3": "其他垃圾/牙签",
"4": "其他垃圾/破碎花盆及碟碗",
"5": "其他垃圾/竹筷",
"6": "厨余垃圾/剩饭剩菜",
"7": "厨余垃圾/大骨头",
"8": "厨余垃圾/水果果皮",
"9": "厨余垃圾/水果果肉",
"10": "厨余垃圾/茶叶渣",
"11": "厨余垃圾/菜叶菜根",
"12": "厨余垃圾/蛋壳",
"13": "厨余垃圾/鱼骨",
"14": "可回收物/充电宝",
"15": "可回收物/包",
"16": "可回收物/化妆品瓶",
"17": "可回收物/塑料玩具",
"18": "可回收物/塑料碗盆",
"19": "可回收物/塑料衣架",
"20": "可回收物/快递纸袋",
"21": "可回收物/插头电线",
"22": "可回收物/旧衣服",
"23": "可回收物/易拉罐",
"24": "可回收物/枕头",
"25": "可回收物/毛绒玩具",
"26": "可回收物/洗发水瓶",
"27": "可回收物/玻璃杯",
"28": "可回收物/皮鞋",
"29": "可回收物/砧板",
"30": "可回收物/纸板箱",
"31": "可回收物/调料瓶",
"32": "可回收物/酒瓶",
"33": "可回收物/金属食品罐",
"34": "可回收物/锅",
"35": "可回收物/食用油桶",
"36": "可回收物/饮料瓶",
"37": "有害垃圾/干电池",
"38": "有害垃圾/软膏",
"39": "有害垃圾/过期药物"
}我这里没有那么多的细分的数据，从数据库中找到了一个四类别的垃圾识别数据集，主要包括：
有害垃圾
可回收垃圾
厨余垃圾
其他垃圾本质都是一样的，都是属于多分类任务。所以这里不需要过于关注数据集类别的问题。
首先看下效果：
基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统
经过一下午的开发，终于完结，项目整体截图如下所示：
其中：
train.py负责整套模型的训练
项目底层使用的是mobilenet系列轻量级的识别模型，对于设备端的压力会更小一点，模型整体体积在8MB左右，一般的PC机在CPU小都是可以轻松推理的。
四种类型垃圾数据截图如下所示：
厨余垃圾：
有害垃圾：
其他垃圾：
可回收垃圾：
results/目录用于存放模型的结果输出，如下所示：
其中evaluation/目录是由evaluation.py模块对测试集进行评估计算得到的评估结果：
我保存了GPU服务器训练时的截图，如下：
红框中的就是我们的模型，可以看到显存占用量不到4GB，还是不错的。
训练日志输出如下所示：
Precision指标计算如下：
Recall指标计算如下：
记录一下，持续学习中，如有需要的话界面顶部有详情，可详细咨询。