基于mobilenet的垃圾分类系统，TensorFlow开发

最近在折腾一个有意思的项目——用手机拍张垃圾照片就能自动分类。试了挺多方案，最后还是用MobileNet搞定了。这货在保持精度的同时体积小得离谱，部署到手机端完全无压力，实测红米Note都能流畅跑起来。

先说说数据的事。Kaggle上有个现成的垃圾分类数据集，四分类（可回收/厨余/有害/其他）约一万五千张图片。处理这种数据最头疼的就是尺寸不统一，直接上OpenCV暴力resize：

train_datagen = ImageDataGenerator(rescale=1./255,
                                   shear_range=0.2,
                                   zoom_range=0.2,
                                   horizontal_flip=True)
train_set = train_datagen.flow_from_directory(
    'dataset/train',
    target_size=(224, 224),
    batch_size=32,
    class_mode='sparse'
)

这里用ImageDataGenerator玩了个小花招，加了随机裁剪和水平翻转。注意class_mode用'sparse'是因为标签是整数形式，比one-hot省内存。实际跑的时候发现有些易拉罐图片会被误判成塑料瓶，于是在数据里混了些商品包装的特写图，效果立竿见影。

基于mobilenet的垃圾分类系统，TensorFlow开发

模型搭建才是重头戏。MobileNetV2的预训练权重真香，但直接拿来用会翻车——原模型是1000类别的ImageNet数据，咱们的垃圾才分4类。这时候需要魔改头部结构：

base_model = tf.keras.applications.MobileNetV2(
    input_shape=(224,224,3),
    include_top=False,
    weights='imagenet'
)
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, activation='relu'),
    tf.keras.layers.Dense(4, activation='softmax')
])

GlobalAveragePooling2D这层大有讲究。原本全连接层参数爆炸，换成全局池化直接砍掉90%参数，精度居然还涨了2个百分点。中间那个Dropout是后来加的——有次训练集准确率飙到99%但验证集卡在85%，明显过拟合了，加个0.3的dropout立马稳住。

训练策略也有门道。初始学习率别超过1e-4，不然分分钟梯度爆炸。这里用了个余弦退火：

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

history = model.fit(
    train_set,
    epochs=30,
    validation_data=val_set,
    callbacks=[tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.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。试了下点外卖剩下的奶茶杯，识别成"其他垃圾"完全正确。不过遇到碎纸机处理的文件碎片还是会误判，可能得加个文本检测模块来辅助，这是后话了。