基于mobilenet的垃圾分类系统TensorFlow开发最近在折腾一个有意思的项目——用手机拍张垃圾照片就能自动分类。试了挺多方案最后还是用MobileNet搞定了。这货在保持精度的同时体积小得离谱部署到手机端完全无压力实测红米Note都能流畅跑起来。先说说数据的事。Kaggle上有个现成的垃圾分类数据集四分类可回收/厨余/有害/其他约一万五千张图片。处理这种数据最头疼的就是尺寸不统一直接上OpenCV暴力resizetrain_datagen ImageDataGenerator(rescale1./255, shear_range0.2, zoom_range0.2, horizontal_flipTrue) train_set train_datagen.flow_from_directory( dataset/train, target_size(224, 224), batch_size32, class_modesparse )这里用ImageDataGenerator玩了个小花招加了随机裁剪和水平翻转。注意class_mode用sparse是因为标签是整数形式比one-hot省内存。实际跑的时候发现有些易拉罐图片会被误判成塑料瓶于是在数据里混了些商品包装的特写图效果立竿见影。基于mobilenet的垃圾分类系统TensorFlow开发模型搭建才是重头戏。MobileNetV2的预训练权重真香但直接拿来用会翻车——原模型是1000类别的ImageNet数据咱们的垃圾才分4类。这时候需要魔改头部结构base_model tf.keras.applications.MobileNetV2( input_shape(224,224,3), include_topFalse, weightsimagenet ) base_model.trainable False # 先冻结预训练层 model tf.keras.Sequential([ base_model, tf.keras.layers.GlobalAveragePooling2D(), tf.keras.layers.Dropout(0.3), tf.keras.layers.Dense(128, activationrelu), tf.keras.layers.Dense(4, activationsoftmax) ])GlobalAveragePooling2D这层大有讲究。原本全连接层参数爆炸换成全局池化直接砍掉90%参数精度居然还涨了2个百分点。中间那个Dropout是后来加的——有次训练集准确率飙到99%但验证集卡在85%明显过拟合了加个0.3的dropout立马稳住。训练策略也有门道。初始学习率别超过1e-4不然分分钟梯度爆炸。这里用了个余弦退火model.compile(optimizertf.keras.optimizers.Adam(learning_rate1e-4), losssparse_categorical_crossentropy, metrics[accuracy]) history model.fit( train_set, epochs30, validation_dataval_set, callbacks[tf.keras.callbacks.ReduceLROnPlateau(monitorval_loss, factor0.2)] )注意验证集loss三次不降就自动砍学习率实测比固定学习率快收敛20个epoch。有个坑得提醒MobileNet的BN层在fine-tune时要小心解冻训练时记得把base_model.trainable设为True并且重新compile才能生效。最后部署时遇到个奇葩问题——安卓端推理结果和PC端不一致。排查发现是图片预处理时RGB通道顺序搞反了加上这个修正才搞定input_array tf.keras.preprocessing.image.img_to_array(bitmap) input_array input_array[:, :, ::-1] # RGB转BGR现在这个模型在麒麟980芯片上跑到60fps毫无压力内存占用不到200MB。试了下点外卖剩下的奶茶杯识别成其他垃圾完全正确。不过遇到碎纸机处理的文件碎片还是会误判可能得加个文本检测模块来辅助这是后话了。