基于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指标计算如下：

 记录一下，持续学习中，如有需要的话界面顶部有详情，可详细咨询。