垃圾分類識(shí)別這件事聽起來不算新鮮但真正動(dòng)手把一套系統(tǒng)從數(shù)據(jù)準(zhǔn)備做到端側(cè)可部署過程中需要踩的坑遠(yuǎn)比想象中多。我這次做的是一套基于TensorFlow的垃圾分類系統(tǒng)核心思路是雙線并行一邊從零訓(xùn)練CNN模型作為基線用來理解卷積網(wǎng)絡(luò)的工作方式和參數(shù)調(diào)優(yōu)另一邊使用MobileNet做遷移學(xué)習(xí)用更少的數(shù)據(jù)拿到更高精度同時(shí)把模型量化為TFLite格式保證它可以跑在手機(jī)、嵌入式開發(fā)板這類資源受限的設(shè)備上。整個(gè)項(xiàng)目做完用攝像頭掃一下就能實(shí)時(shí)輸出垃圾類別是一個(gè)邏輯完整、可以繼續(xù)擴(kuò)展的工程實(shí)踐。這篇文章適合剛?cè)腴T深度學(xué)習(xí)、想做一個(gè)真實(shí)項(xiàng)目練手的同學(xué)也適合正在做課程設(shè)計(jì)、畢業(yè)設(shè)計(jì)或產(chǎn)品原型的開發(fā)者。我會(huì)把每一個(gè)關(guān)鍵選擇的“為什么”講清楚不只是貼代碼還會(huì)分享我在實(shí)驗(yàn)里遇到的高頻坑和解決辦法。你能從里面直接拿走一套可復(fù)現(xiàn)的方案。1. 先把垃圾分類系統(tǒng)拆清楚需求、方案與技術(shù)棧1.1 需求拆解你要做的不是一個(gè)“模型”而是一套系統(tǒng)很多人在拿到“垃圾分類”這個(gè)題目后第一反應(yīng)就是找數(shù)據(jù)集、跑分類網(wǎng)絡(luò)但真正落地時(shí)要回答的問題遠(yuǎn)不止這些。第一個(gè)問題是“輸入是什么”是單張圖片還是攝像頭視頻流手機(jī)拍照識(shí)別和實(shí)時(shí)視頻識(shí)別對(duì)模型大小和推理速度的要求完全不同。第二個(gè)問題是“輸出是什么”是輸出一個(gè)類別標(biāo)簽還是需要在界面上同時(shí)顯示置信度和對(duì)應(yīng)垃圾的投放規(guī)則第三個(gè)問題是“運(yùn)行環(huán)境”是在電腦的GPU上演示還是要部署到無GPU的移動(dòng)端我給自己定的目標(biāo)很實(shí)際輸入是普通攝像頭實(shí)時(shí)畫面或照片輸出是垃圾類別和置信度最終模型可以轉(zhuǎn)換成TFLite格式在CPU設(shè)備上流暢運(yùn)行。從這個(gè)目標(biāo)倒推模型就不能選那種追求極致精度但體積好幾GB的超大網(wǎng)絡(luò)訓(xùn)練框架則定在TensorFlow 2.x因?yàn)樗鼜挠?xùn)練到部署的鏈路最完整對(duì)新手也友好。1.2 為什么同時(shí)用自訓(xùn)練CNN和MobileNet遷移學(xué)習(xí)兩條路線如果只追求結(jié)果直接用現(xiàn)成的預(yù)訓(xùn)練模型做遷移學(xué)習(xí)就可以了不需要自己從零構(gòu)建CNN。但我還是堅(jiān)持把兩條路線都做出來原因有兩點(diǎn)。第一自訓(xùn)練CNN是理解深度學(xué)習(xí)底層原理的最佳路徑。通過自己搭建卷積層、池化層、全連接層你能直觀感受到卷積核怎么提特征、池化怎么降維、Dropout怎么抑制過擬合。這些感知在你后續(xù)調(diào)整任何模型時(shí)都非常有用。第二在實(shí)際數(shù)據(jù)量有限的情況下從零訓(xùn)練CNN很難追上遷移學(xué)習(xí)的效果因?yàn)槔诸悢?shù)據(jù)集的規(guī)模和復(fù)雜度都遠(yuǎn)不如大規(guī)模自然圖像數(shù)據(jù)集模型很難只靠少量樣本學(xué)到足夠泛化的特征。所以我的策略是把自訓(xùn)練CNN當(dāng)作基線模型用來建立“最差效果”認(rèn)知把MobileNet遷移學(xué)習(xí)作為正式方案負(fù)責(zé)最終落地精度和部署可行性。1.3 技術(shù)棧與總體流程設(shè)計(jì)整個(gè)系統(tǒng)的技術(shù)選型如下表選型邏輯是“優(yōu)先考慮完整性和可復(fù)現(xiàn)性”。模塊選型選擇理由深度學(xué)習(xí)框架TensorFlow 2.x訓(xùn)練接口成熟TFLite轉(zhuǎn)換方便部署鏈路完整基礎(chǔ)模型自訓(xùn)練CNN MobileNetV2兼顧原理學(xué)習(xí)與落地性能MobileNet在移動(dòng)端友好圖像處理OpenCV、TensorFlow數(shù)據(jù)管道攝像頭讀取、圖像預(yù)處理和增強(qiáng)都能覆蓋訓(xùn)練平臺(tái)CPU可跑GPU更好數(shù)據(jù)量不大時(shí)CPU也能完成但GPU能顯著加快迭代部署格式TFLite模型小、推理快不依賴完整TensorFlow環(huán)境整體流程分成四段數(shù)據(jù)準(zhǔn)備與預(yù)處理、自訓(xùn)練CNN做基線、MobileNet遷移學(xué)習(xí)調(diào)優(yōu)、TFLite導(dǎo)出與端側(cè)驗(yàn)證。每一段都不是孤立的數(shù)據(jù)準(zhǔn)備的質(zhì)量直接影響后面所有環(huán)節(jié)所以我建議不要在數(shù)據(jù)階段節(jié)省時(shí)間。2. 數(shù)據(jù)準(zhǔn)備與預(yù)處理實(shí)戰(zhàn)2.1 數(shù)據(jù)來源與類別設(shè)計(jì)先定類別再談效果垃圾分類系統(tǒng)的第一步不是寫模型而是確定類別體系。常見的公共垃圾分類數(shù)據(jù)集多以“可回收垃圾、廚余垃圾、有害垃圾、其他垃圾”四大類為頂層再細(xì)化到數(shù)十個(gè)子類。我建議第一次做項(xiàng)目的人不要一上來就挑戰(zhàn)上百個(gè)細(xì)分類別先做4到10類等整個(gè)訓(xùn)練部署流程跑通之后再逐步擴(kuò)充。我是從“紙板、玻璃瓶、金屬罐、塑料瓶、塑料袋、果皮、廢電池、舊衣服”這8類開始的每類收集600到1000張圖片。如果你使用公開數(shù)據(jù)集一定要先做一次人工抽檢因?yàn)檫@類數(shù)據(jù)集中最容易出現(xiàn)的問題是圖片分辨率參差不齊、存在重復(fù)圖片、少數(shù)圖片標(biāo)注明顯錯(cuò)誤比如把塑料瓶標(biāo)成玻璃瓶。如果模型學(xué)習(xí)到了錯(cuò)標(biāo)特征后面的精確率再高也是虛的。數(shù)據(jù)不足的部分可以自己拍攝補(bǔ)充注意覆蓋不同角度、不同光線和不同背景這樣模型才不至于只在某一類固定背景上表現(xiàn)好。2.2 預(yù)處理細(xì)節(jié)尺寸、歸一化與通道順序預(yù)處理看起來簡單卻決定了模型能否正確接收?qǐng)D片。首先必須把所有圖片統(tǒng)一成相同尺寸否則CNN全連接層的輸入維度會(huì)不一致訓(xùn)練直接報(bào)錯(cuò)。MobileNetV2的默認(rèn)輸入尺寸是224x224自訓(xùn)練CNN則根據(jù)計(jì)算資源可以選擇96x96或128x128。其次把像素值從0到255歸一化。使用ImageDataGenerator時(shí)設(shè)置rescale1./255即可使用MobileNetV2時(shí)要注意它有自己的預(yù)處理函數(shù)不能簡單除以255就算完。還要注意通道順序。TensorFlow默認(rèn)是“高度、寬度、通道”的NHWC格式一般讀圖后不做特殊處理即可。如果你習(xí)慣用OpenCV讀圖要記得OpenCV讀取的是BGR順序需要轉(zhuǎn)換為RGB否則模型看到的顏色通道是反的訓(xùn)練效果會(huì)明顯變差。2.3 數(shù)據(jù)增強(qiáng)讓模型在少量數(shù)據(jù)里“見多識(shí)廣”數(shù)據(jù)增強(qiáng)不是單純?yōu)榱嗽黾訄D片數(shù)量而是讓模型對(duì)拍攝時(shí)的真實(shí)變化不敏感。我使用的增強(qiáng)方式包括隨機(jī)旋轉(zhuǎn)、水平翻轉(zhuǎn)、寬度和高度平移、縮放、剪切變換、亮度擾動(dòng)。用TensorFlow的ImageDataGenerator實(shí)現(xiàn)只需要幾行配置。from tensorflow.keras.preprocessing.image import ImageDataGenerator train_datagen ImageDataGenerator( rescale1./255, rotation_range20, width_shift_range0.2, height_shift_range0.2, shear_range0.2, zoom_range0.2, horizontal_flipTrue, brightness_range[0.8, 1.2], fill_modenearest )這些增強(qiáng)參數(shù)也不是越大越好。比如旋轉(zhuǎn)角度超過30度很多垃圾對(duì)象的類別特征可能被破壞反而加大學(xué)習(xí)難度。我建議先用輕中度增強(qiáng)觀察驗(yàn)證集準(zhǔn)確率是否提升再逐步調(diào)整。務(wù)必記住驗(yàn)證集和測試集不能使用與訓(xùn)練集相同的數(shù)據(jù)增強(qiáng)最多只做縮放和歸一化否則驗(yàn)證結(jié)果會(huì)虛高不能反映模型真實(shí)泛化能力。2.4 數(shù)據(jù)集劃分與目錄組織數(shù)據(jù)劃分要嚴(yán)格做到分層抽樣避免某一類圖片全部集中在訓(xùn)練集而驗(yàn)證集只剩下其他類別。對(duì)多分類任務(wù)最簡單有效的做法是使用train_test_split并指定stratifyy。import os import shutil from sklearn.model_selection import train_test_split all_images [] for category in categories: imgs [os.path.join(category, f) for f in os.listdir(category)] all_images.extend(imgs) labels [category] * len(imgs) train_imgs, val_imgs, train_labels, val_labels train_test_split( all_images, labels, test_size0.2, stratifylabels, random_state42 )目錄結(jié)構(gòu)建議采用規(guī)范的分類文件夾方式data/ ├── train/ │ ├── cardboard/ │ ├── metal/ │ └── ... ├── val/ └── test/直接把圖片按類別放進(jìn)對(duì)應(yīng)的文件夾后續(xù)用ImageDataGenerator的flow_from_directory加載非常方便也方便人工檢查數(shù)據(jù)。我習(xí)慣把訓(xùn)練、驗(yàn)證、測試按7:2:1或8:1:1劃分測試集留到最后才用避免在調(diào)試過程中反復(fù)看測試集結(jié)果產(chǎn)生隱性的“測試集過擬合”。3. 從零訓(xùn)練CNN網(wǎng)絡(luò)結(jié)構(gòu)、參數(shù)計(jì)算與訓(xùn)練回調(diào)3.1 CNN為什么能處理圖片從卷積核到特征圖CNN的核心假設(shè)是圖像具有局部相關(guān)性。一個(gè)像素往往與周圍的像素關(guān)系更密切而遠(yuǎn)處像素關(guān)系較弱。卷積層通過一組可學(xué)習(xí)的卷積核在圖片上滑動(dòng)每個(gè)卷積核負(fù)責(zé)檢測某一種局部特征比如邊緣、顏色塊、紋理。多個(gè)卷積核疊加后淺層網(wǎng)絡(luò)學(xué)到的是線條和角點(diǎn)深層網(wǎng)絡(luò)學(xué)到的是“瓶口”“把手”這類語義部件最后用全連接層把這些特征組合成類別判斷。池化層的作用是在保留主要特征的同時(shí)縮小特征圖尺寸既減少計(jì)算量也增強(qiáng)平移不變性。整個(gè)過程很像我們?nèi)搜巯瓤淳植考?xì)節(jié)再綜合判斷整體類別。理解了這一點(diǎn)你就能明白為什么那些把卷積核堆得很深卻不做歸一化的模型經(jīng)常出現(xiàn)訓(xùn)練震蕩。3.2 網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計(jì)不只是“越深越好”我設(shè)計(jì)的自訓(xùn)練CNN結(jié)構(gòu)如下適合8個(gè)分類類別輸入尺寸為224x224x3卷積層132個(gè)3x3卷積核BatchNormalizationReLU2x2最大池化卷積層264個(gè)3x3卷積核BatchNormalizationReLU2x2最大池化卷積層3128個(gè)3x3卷積核BatchNormalizationReLU全局平均池化全連接層128個(gè)神經(jīng)元ReLUDropout 0.5輸出層8個(gè)神經(jīng)元Softmax為什么第三層用全局平均池化而不是繼續(xù)展平因?yàn)槿绻苯覨latten前面特征圖的尺寸還會(huì)很大參數(shù)量會(huì)急劇增加。比如128個(gè)特征圖每個(gè)尺寸為28x28展平后就是100352個(gè)輸入節(jié)點(diǎn)全連接層參數(shù)量巨大非常容易過擬合。全局平均池化把每個(gè)特征圖變成一個(gè)值參數(shù)量大幅下降效果在大多數(shù)分類任務(wù)上不輸Flatten。實(shí)際參數(shù)量估算第一層卷積是3x3x3x32再加32個(gè)偏置一共896個(gè)參數(shù)第二層3x3x32x64再加64個(gè)偏置約18496個(gè)參數(shù)第三層3x3x64x128再加128約73856個(gè)參數(shù)。整個(gè)模型參數(shù)量不到幾十萬相比于幾百萬甚至上億參數(shù)的大模型已經(jīng)非常輕量。3.3 TensorFlow代碼實(shí)現(xiàn)import tensorflow as tf from tensorflow.keras import layers, models def build_cnn(input_shape(224, 224, 3), num_classes8): model models.Sequential([ layers.Conv2D(32, (3, 3), activationrelu, input_shapeinput_shape), layers.BatchNormalization(), layers.MaxPooling2D(pool_size(2, 2)), layers.Conv2D(64, (3, 3), activationrelu), layers.BatchNormalization(), layers.MaxPooling2D(pool_size(2, 2)), layers.Conv2D(128, (3, 3), activationrelu), layers.BatchNormalization(), layers.GlobalAveragePooling2D(), layers.Dense(128, activationrelu), layers.Dropout(0.5), layers.Dense(num_classes, activationsoftmax) ]) return model這里有兩個(gè)容易被忽略的細(xì)節(jié)。第一是BatchNormalization放在卷積和激活之間很多新手喜歡先激活再歸一化效果往往會(huì)差一點(diǎn)。第二是卷積層后面接了池化不要急著堆很多卷積層因?yàn)樾?shù)據(jù)集上的深度神經(jīng)網(wǎng)絡(luò)很容易過擬合淺一點(diǎn)反而好訓(xùn)練。3.4 訓(xùn)練配置與回調(diào)機(jī)制訓(xùn)練配置決定了模型收斂速度和最終效果不要隨便選參數(shù)。model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-3), losscategorical_crossentropy, metrics[accuracy] ) callbacks [ tf.keras.callbacks.EarlyStopping(monitorval_loss, patience10, restore_best_weightsTrue), tf.keras.callbacks.ReduceLROnPlateau(monitorval_loss, factor0.5, patience5), tf.keras.callbacks.ModelCheckpoint(cnn_best.h5, monitorval_accuracy, save_best_onlyTrue) ] history model.fit( train_generator, validation_dataval_generator, epochs50, callbackscallbacks )categorical_crossentropy要求標(biāo)簽是one-hot編碼flow_from_directory默認(rèn)會(huì)用categorical模式輸出one-hot標(biāo)簽所以這里正好匹配。EarlyStopping和ReduceLROnPlateau是好搭檔當(dāng)驗(yàn)證集loss不再下降時(shí)學(xué)習(xí)率會(huì)自動(dòng)降一半給模型一個(gè)更精細(xì)的收斂機(jī)會(huì)如果連續(xù)10個(gè)epoch仍然沒有改善就提前停止避免浪費(fèi)時(shí)間。3.5 評(píng)估模型準(zhǔn)確率不夠要看混淆矩陣訓(xùn)練結(jié)束后不要只看總準(zhǔn)確率一定要看分類報(bào)告。8個(gè)類別的垃圾中可能“塑料瓶”和“玻璃瓶”經(jīng)常混淆而其他類別都很好這時(shí)候總準(zhǔn)確率可能顯示92%看起來不錯(cuò)但特定兩類完全不可用。正確做法是在測試集上輸出混淆矩陣和逐類精確率、召回率。from sklearn.metrics import classification_report, confusion_matrix import numpy as np y_pred model.predict(test_generator) y_pred_classes np.argmax(y_pred, axis1) y_true test_generator.classes print(classification_report(y_true, y_pred_classes, target_namescategories)) print(confusion_matrix(y_true, y_pred_classes))一個(gè)真實(shí)項(xiàng)目中如果“塑料袋”因?yàn)樘该?、反光?yán)重而被大量誤判為“塑料瓶”說明模型還沒有把“形狀可變形”這個(gè)特征學(xué)會(huì)。這時(shí)候優(yōu)先補(bǔ)充更多透明的、揉成一團(tuán)的塑料袋圖片比盲目加深網(wǎng)絡(luò)更有效。這也是小項(xiàng)目里“數(shù)據(jù)比模型更值錢”的典型例子。4. MobileNet遷移學(xué)習(xí)實(shí)戰(zhàn)少樣本也能拿到高精度4.1 MobileNet的核心優(yōu)勢輕量而高效MobileNet不是某一種固定結(jié)構(gòu)而是一系列為移動(dòng)端設(shè)計(jì)的輕量網(wǎng)絡(luò)。它的核心創(chuàng)新是深度可分離卷積把普通卷積拆分成“逐通道卷積”和“逐點(diǎn)卷積”。簡單理解普通卷積一次性同時(shí)處理空間信息和通道信息而深度可分離卷積先對(duì)每個(gè)通道做空間卷積再用1x1卷積去融合通道信息。這樣計(jì)算量可以下降一個(gè)數(shù)量級(jí)而精度損失相對(duì)可控。對(duì)于垃圾分類這個(gè)場景模型是要部署到普通設(shè)備上的MobileNet的輕量特性直接決定了“可用”。我選擇MobileNetV2而不是MobileNetV1是因?yàn)閂2引入了線性瓶頸和倒殘差結(jié)構(gòu)在相同體積下精度更高很多公開部署案例都驗(yàn)證過它的穩(wěn)定性。4.2 遷移學(xué)習(xí)為什么要借大規(guī)模預(yù)訓(xùn)練模型的“眼力”垃圾分類圖片雖然有自己的特點(diǎn)但底層視覺特征比如邊緣、紋理、形狀、顏色分布與大規(guī)模自然圖像數(shù)據(jù)集中學(xué)到的特征非常相似。預(yù)訓(xùn)練模型已經(jīng)在海量圖片上學(xué)會(huì)了識(shí)別這些通用特征我們不需要從零開始教它“什么是邊緣”。遷移學(xué)習(xí)有兩種常見策略。第一種是特征提取把預(yù)訓(xùn)練模型全部凍結(jié)只修改最后幾層適合每類只有幾百張圖片的小數(shù)據(jù)集。第二種是微調(diào)先凍結(jié)預(yù)訓(xùn)練模型訓(xùn)練分類頭再解凍部分高層卷積層用很低的學(xué)習(xí)率繼續(xù)訓(xùn)練適合數(shù)據(jù)量稍多且希望進(jìn)一步提升精度的場景。我用一張表說明區(qū)別策略訓(xùn)練時(shí)間所需數(shù)據(jù)量精度上限適合場景特征提取短每類200~500張中等偏高快速搭建基線、CPU訓(xùn)練微調(diào)較長每類500張以上更高追求最終部署精度4.3 特征提取模式實(shí)現(xiàn)先凍結(jié)全部卷積基座下面這段代碼創(chuàng)建了MobileNetV2遷移學(xué)習(xí)模型關(guān)鍵點(diǎn)在于base_model.trainable False。from tensorflow.keras.applications import MobileNetV2 from tensorflow.keras.applications.mobilenet_v2 import preprocess_input from tensorflow.keras import layers, models base_model MobileNetV2( weightsimagenet, include_topFalse, input_shape(224, 224, 3) ) base_model.trainable False inputs tf.keras.Input(shape(224, 224, 3)) x preprocess_input(inputs) x base_model(x, trainingFalse) x layers.GlobalAveragePooling2D()(x) x layers.Dropout(0.2)(x) outputs layers.Dense(num_classes, activationsoftmax)(x) model models.Model(inputs, outputs) model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-3), losscategorical_crossentropy, metrics[accuracy] )注意這里的preprocess_input是MobileNetV2專用的預(yù)處理方式不是簡單除以255。把輸入先做預(yù)處理再進(jìn)入模型才能和預(yù)訓(xùn)練時(shí)看到的數(shù)據(jù)分布保持一致。還有一個(gè)容易被忽略的細(xì)節(jié)因?yàn)閎ase_model被凍結(jié)在推理模式下trainingFalseBatchNormalization層會(huì)使用移動(dòng)平均而不是當(dāng)前批次統(tǒng)計(jì)量這樣能避免因?yàn)閮鼋Y(jié)BN層導(dǎo)致的訓(xùn)練不穩(wěn)定。4.4 微調(diào)模式實(shí)現(xiàn)小學(xué)習(xí)率解凍高層特征提取模式訓(xùn)練幾個(gè)epoch后如果驗(yàn)證準(zhǔn)確率進(jìn)入了平臺(tái)期就可以嘗試微調(diào)。微調(diào)的思路不是一次性把整個(gè)網(wǎng)絡(luò)都解凍那樣學(xué)習(xí)率稍大一點(diǎn)就會(huì)讓已經(jīng)學(xué)好的預(yù)訓(xùn)練特征被破壞。我的做法是先解凍整個(gè)MobileNetV2然后手動(dòng)把比較淺的層繼續(xù)凍結(jié)只讓最后幾十層跟著訓(xùn)練。base_model.trainable True # 具體凍結(jié)層數(shù)需要根據(jù)模型結(jié)構(gòu)調(diào)整一般保留前100層左右不動(dòng) for layer in base_model.layers[:100]: layer.trainable False model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-5), losscategorical_crossentropy, metrics[accuracy] )這個(gè)階段的學(xué)習(xí)率一定要比前面低兩個(gè)數(shù)量級(jí)以上。我用過1e-5也用過5e-6整體都很穩(wěn)定。如果學(xué)習(xí)率突然拉高到1e-4很快就看到loss沖高然后模型再也回不到原來的狀態(tài)這就是很多同學(xué)遇到的“微調(diào)直接廢掉”的原因。4.5 小樣本訓(xùn)練的幾個(gè)關(guān)鍵技巧在每類只有幾百張圖片時(shí)有幾個(gè)技巧非常管用。第一是先用特征提取模式得到一個(gè)不錯(cuò)的初始化點(diǎn)再進(jìn)入微調(diào)不要在一開始就解凍全部層。第二是在分類頭之前加Dropout層雖然預(yù)訓(xùn)練模型已經(jīng)很強(qiáng)但在小數(shù)據(jù)集上依然很容易過擬合。第三是合理使用早停讓訓(xùn)練在驗(yàn)證集開始下降時(shí)立刻停止。第四是如果某類圖片數(shù)量特別少可以為這類圖片做更多增強(qiáng)或者用類別權(quán)重懲罰模型對(duì)少數(shù)類的“忽視”。我實(shí)測下來MobileNet遷移學(xué)習(xí)在每類500張左右圖片的情況下驗(yàn)證準(zhǔn)確率比自訓(xùn)練CNN高出8到12個(gè)百分點(diǎn)而且模型體積只有約十幾MB轉(zhuǎn)換TFLite之后還能再壓縮一半以上。這也是為什么最終方案我毫不猶豫選擇遷移學(xué)習(xí)的原因。5. 讓模型真正可用TFLite導(dǎo)出與端側(cè)推理5.1 部署目標(biāo)模型要“變小變快”而不是只追求準(zhǔn)確率很多同學(xué)訓(xùn)練完模型保存一個(gè).h5文件就認(rèn)為項(xiàng)目完成了。但對(duì)“可用”來說這遠(yuǎn)遠(yuǎn)不夠。h5或SavedModel格式需要完整的TensorFlow環(huán)境才能加載CPU設(shè)備上跑起來還特別慢。手機(jī)、嵌入式板卡根本沒有辦法直接運(yùn)行大模型。所以我的部署目標(biāo)是轉(zhuǎn)換成本低的TFLite格式。TFLite是TensorFlow為移動(dòng)端和嵌入式設(shè)備設(shè)計(jì)的推理格式它執(zhí)行時(shí)不需要完整訓(xùn)練框架模型體積更小推理速度快。轉(zhuǎn)換時(shí)還可以做量化把模型中的權(quán)重從32位浮點(diǎn)數(shù)變成16位浮點(diǎn)數(shù)甚至8位整數(shù)進(jìn)一步壓縮體積。5.2 從SavedModel到TFLite的轉(zhuǎn)換操作先訓(xùn)練并保存模型推薦保存為SavedModel格式因?yàn)樗鼘?duì)部署鏈路兼容性最好。model.save(garbage_mobilenet_savedmodel)然后使用TFLite轉(zhuǎn)換器。最簡單的轉(zhuǎn)換如下import tensorflow as tf converter tf.lite.TFLiteConverter.from_saved_model(garbage_mobilenet_savedmodel) converter.optimizations [tf.lite.Optimize.DEFAULT] tflite_model converter.convert() with open(garbage_mobilenet.tflite, wb) as f: f.write(tflite_model)這里tf.lite.Optimize.DEFAULT會(huì)盡可能壓縮模型。由于訓(xùn)練時(shí)模型權(quán)重分布比較合理量化后精度損失通常很小。如果后續(xù)實(shí)測發(fā)現(xiàn)精度掉得比較多可以改用僅對(duì)權(quán)重做量化而不對(duì)激活值做量化或者干脆不量化只保留FP32模型根據(jù)自己的硬件算力做取舍。5.3 使用TFLite解釋器在Python端推理TFLite模型不能直接喂給原來的模型預(yù)測函數(shù)需要用TFLite解釋器加載。我寫了一個(gè)可復(fù)用的推理函數(shù)import numpy as np import tensorflow as tf interpreter tf.lite.Interpreter(model_pathgarbage_mobilenet.tflite) interpreter.allocate_tensors() input_details interpreter.get_input_details() output_details interpreter.get_output_details() def predict_tflite(image): image cv2.resize(image, (224, 224)) image image.astype(np.float32) image preprocess_input(image) # 和訓(xùn)練時(shí)保持一致 image np.expand_dims(image, axis0) interpreter.set_tensor(input_details[0][index], image) interpreter.invoke() output interpreter.get_tensor(output_details[0][index]) return output部署時(shí)最容易出的問題就是預(yù)處理不一致。訓(xùn)練時(shí)你用了MobileNetV2的preprocess_input推理時(shí)卻用了簡單除以255模型輸出的置信度會(huì)很奇怪。如果你自己在模型里加了歸一化層那么推理時(shí)就不需要再做一遍。在寫部署代碼之前先花十分鐘核對(duì)接下來的每一步比訓(xùn)練時(shí)調(diào)參還要重要。5.4 精度與速度的取舍量化不是越高越好我實(shí)際對(duì)比過三份模型原始FP32、FP16半精度、INT8量化。模型模型體積推理速度準(zhǔn)確率變化FP32約14MB基準(zhǔn)基準(zhǔn)FP16約7MB更快基本不變INT8約3.5MB最快可能下降1%~3%INT8量化在小模型上通常不會(huì)有太大損失但如果你訓(xùn)練時(shí)的數(shù)據(jù)分布與部署現(xiàn)場差異較大掉點(diǎn)就會(huì)更明顯。一個(gè)穩(wěn)妥的做法是先在開發(fā)板上跑一遍量化模型如果準(zhǔn)確率下降超過預(yù)期就退回FP16。端側(cè)推理還可以設(shè)置線程數(shù)比如在雙核處理器上嘗試interpreter.set_num_threads(2)有時(shí)候速度提升比換模型更直接。6. 常見問題與高頻踩坑實(shí)錄6.1 訓(xùn)練集較小模型反復(fù)過擬合典型現(xiàn)象是訓(xùn)練準(zhǔn)確率很快升到98%驗(yàn)證準(zhǔn)確率卻停在70%。這通常意味著模型把訓(xùn)練數(shù)據(jù)死記硬背下來了沒有真正學(xué)會(huì)泛化。解決順序分三步先加強(qiáng)數(shù)據(jù)增強(qiáng)增加旋轉(zhuǎn)、翻轉(zhuǎn)、色彩抖動(dòng)的強(qiáng)度再增強(qiáng)模型正則化增大Dropout比例或減少全連接層神經(jīng)元最后如果還是不行就換遷移學(xué)習(xí)模型而不是繼續(xù)堆CNN層數(shù)。對(duì)小數(shù)據(jù)來說換網(wǎng)絡(luò)帶來的收益往往大于調(diào)參。6.2 類別不均衡導(dǎo)致模型“永遠(yuǎn)猜多數(shù)類”如果“塑料瓶”圖片有3000張“舊衣服”只有200張模型很可能為了整體準(zhǔn)確率把舊衣服全部忽略。在訓(xùn)練損失上體現(xiàn)為少數(shù)類幾乎不被預(yù)測。對(duì)策有兩個(gè)一是收集更多少數(shù)類圖片這是最根本的二是在class_weight里給少數(shù)類更大權(quán)重。TensorFlow訓(xùn)練時(shí)可以直接傳入class_weight參數(shù)不同框架語法略有差異核心思想是讓少數(shù)類樣本被分錯(cuò)時(shí)產(chǎn)生更大的損失。6.3 驗(yàn)證集準(zhǔn)確率高測試集卻崩了這是最隱蔽的坑。常見原因是數(shù)據(jù)泄漏同一個(gè)物體的多張不同角度照片被同時(shí)分到了訓(xùn)練集和驗(yàn)證集導(dǎo)致驗(yàn)證集“見過”訓(xùn)練數(shù)據(jù)。另外數(shù)據(jù)增強(qiáng)如果被錯(cuò)誤地應(yīng)用到了驗(yàn)證集也會(huì)讓驗(yàn)證分?jǐn)?shù)虛高。我建議劃分?jǐn)?shù)據(jù)集時(shí)按“物體實(shí)例”而不是按“單張圖片”去劃分比如同一組連續(xù)拍攝的塑料瓶圖片全部劃到同一個(gè)集合不要打散。這樣做更接近真實(shí)場景。6.4 遷移學(xué)習(xí)微調(diào)時(shí)loss突然沖高這是我在遷移學(xué)習(xí)里踩過最狠的坑。原因通常是解凍后學(xué)習(xí)率設(shè)置過大導(dǎo)致預(yù)訓(xùn)練權(quán)重被快速破壞也可能是模型里BatchNormalization層在trainingTrue時(shí)重新統(tǒng)計(jì)了少量樣本的均值方差而樣本量太小統(tǒng)計(jì)不穩(wěn)定。解決方法是把學(xué)習(xí)率降到1e-5以下同時(shí)確保微調(diào)階段的batch size不要太小盡量保持在32以上。如果loss還是沖高重新從特征提取階段的模型恢復(fù)再降低解凍層數(shù)。6.5 端側(cè)部署后推理速度達(dá)不到實(shí)時(shí)如果攝像頭畫面看上去有明顯的卡頓先檢查模型輸入分辨率是不是太高。比如你訓(xùn)練時(shí)用224x224部署時(shí)卻把原圖先縮放到640x640再送進(jìn)模型速度自然慢。輸入超過模型需要的大小并不會(huì)提升精度。其次檢查量化是否生效先用interpreter打印模型輸入數(shù)據(jù)類型如果還是float32就確認(rèn)轉(zhuǎn)換時(shí)是否真的加上了optimizations。兩條都排查完基本可以解決問題。7. 寫在最后一些真實(shí)體會(huì)7.1 復(fù)現(xiàn)這個(gè)項(xiàng)目的建議順序我自己跑這個(gè)項(xiàng)目時(shí)最受益的是把“自訓(xùn)練CNN”和“MobileNet遷移學(xué)習(xí)”放在同一套數(shù)據(jù)上對(duì)比。這樣你能直觀感受到數(shù)據(jù)量、模型容量和泛化能力之間的關(guān)系而不是單純迷信某個(gè)高大上的網(wǎng)絡(luò)。建議復(fù)現(xiàn)時(shí)也按這個(gè)順序來先準(zhǔn)備數(shù)據(jù)再跑一遍CNN基線然后用MobileNet遷移學(xué)習(xí)做正式方案最后再做TFLite量化部署。不要急著跳過基線直接做遷移學(xué)習(xí)少了那一步很多訓(xùn)練現(xiàn)象你很難解釋。7.2 這個(gè)項(xiàng)目后續(xù)還能往哪里擴(kuò)展完成第一版后可以做的擴(kuò)展方向很多。比如給系統(tǒng)加入目標(biāo)檢測前置模塊先用檢測模型定位畫面中的垃圾區(qū)域再裁剪送入分類模型這樣就能處理畫面里有多個(gè)垃圾對(duì)象的情況。也可以把單幀圖片分類改成視頻流分類加入時(shí)序平滑策略避免連續(xù)幀分類結(jié)果抖動(dòng)太厲害。再進(jìn)一步還可以在部署腳本里加入簡單的置信度閾值邏輯置信度低時(shí)提示用戶重新拍攝或反饋人工標(biāo)注形成一個(gè)小閉環(huán)。垃圾分類系統(tǒng)并不是什么“高不可攀”的項(xiàng)目它真正考驗(yàn)的是你把數(shù)據(jù)、模型、部署串聯(lián)起來的工程能力。希望這篇記錄能幫你少走幾步彎路。