【導讀】TensorFlow.js 的發布可以說是 JS 社區開發者的福音！但是在瀏覽器中訓練一些模型還是會存在一些問題與不同，如何可以讓訓練效果更好？本文的作者，是一位前端工程師，經過自己不斷的經驗積累，為大家總結了 18 個 Tips，希望可以幫助大家訓練出更好的模型。

TensorFlow.js 發布之后我就把之前訓練的目標/人臉檢測和人臉識別的模型往 TensorFlow.js 里導，我發現有些模型在瀏覽器里運行的效果還相當不錯。感覺 TensorFlow.js 讓我們搞前端的也潮了一把。

雖說瀏覽器也能跑深度學習模型了，這些模型終歸不是為在瀏覽器里運行設計的，所以很多限制和挑戰也就隨之而來了。就拿目標檢測來說，不說實時檢測，就是維持一定的幀率恐怕都很困難。更別提動輒上百兆的模型給用戶瀏覽器和帶寬（手機端的話）帶來的壓力了。

不過只要我們遵循一定的原則，用卷積神經網絡 CNN 和 TensorFlow.js 在瀏覽器里訓練個像樣的深度學習模型并非癡人說夢。從下面圖里可以看到，我訓練的這幾個模型大小都控制在了 2 MB 以下，最小的才 3 KB。

大家可能心中會有個疑問：你腦殘嗎？要用瀏覽器訓練模型？對，用自己電腦、服務器、集群或者云來訓練深度學習模型肯定是一條正道，但并非人人都有錢用NVIDIA GTX 1080 Ti 或者Titan X（尤其是顯卡集體大漲價之后）。這時，在瀏覽器中訓練深度學習模型的優勢就體現出來了，有了 WebGL 和 TensorFLow.js 我用電腦上的 AMD GPU 也能很方便地訓練深度學習模型。

對目標識別問題，為了穩妥起見通常都會建議大家用一些現成的架構比如YOLO、SSD、殘差網絡 ResNet 或 MobileNet ，但我個人認為如果完全照搬的話，在瀏覽器上訓練效果肯定是不好的。在瀏覽器上訓練就要求模型要小、要快、要越容易訓練越好。下面我們就從模型架構、訓練和調試等幾個方面來看看如何才能做到這三點。

模型架構

▌1. 控制模型大小

控制模型的規模很重要。如果模型架構太大太復雜，訓練和運行的速度都會降低，從瀏覽器載入模型度速度也會變慢。控制模型的規模說起來簡單，難的是取得準確率和模型規模之間的平衡。如果準確率達不到要求，模型再小也是廢物。

▌2. 使用深度可分離卷積操作

與標準卷積操作不同，深度可分離卷積先對每個通道進行卷積操作，之后再進行1X1跨通道卷積。這樣做的好處是可以大大減小參數個數，所以模型運行速度會有很大提升，資源的消耗和訓練速度也會有所提升。深度可分離卷積操作的過程如下圖所示：

MobileNet 和 Xception 都使用了深度可分離卷積，TensorFlow.js 版本的 MobileNet 和 PoseNet 中你也能見到深度可分離卷積的身影。雖然深度可分離卷積對模型準確率的影響還有爭議，但從我個人的經驗來看在瀏覽器里訓練模型用它肯定沒錯。

第一層我推薦用標準的 conv2d 操作來保持提取完特征的通道之間的關系。因為第一層一般參數不多，所以對性能的影響不大。

其他卷積層就可以都用深度可分離卷積了。比如這里我們就使用了兩個過濾器。

這里 tf.separableConv2d 使用的卷積核結構分別是[3,3,32,1]和[1,1,32,64]。

▌3.運用跳躍連接和密集塊

隨著網絡層數的增加，梯度消失問題出現的可能性也會增大。梯度消失會造成損失函數下降太慢訓練時間超長或者干脆失敗。ResNet 和 DenseNet 中采用的跳躍連接則能避免這一問題。簡單說來跳躍連接就是把某些層的輸出跳過激活函數直接傳給網絡深處的隱藏層作為輸入，如下圖所示：

這樣就避免了因為激活函數和鏈式求導造成的梯度消失問題，我們也能根據需求增加網絡的層數了。

顯然跳躍連接隱含的一個要求就是連接的兩層輸出和輸入的格式必須能對應得上。我們要用殘差網絡的話，那最好保證兩層的過濾器數目和填充都一致而且步幅為1（不過肯定有其它做法來保證格式對應）。

一開始我模仿殘差網絡的思路隔一層加一個跳躍連接（如下圖）。不過我發現密集塊效果更好，模型收斂的速度比加跳躍連接快得多。

下面我們就來看看具體的代碼，這里的密集塊有四個深度可分離卷積層，其中第一層我把步幅設為 2 來改變輸入的大小。

▌4.激活函數選ReLU

在瀏覽器里訓練深度網絡的話激活函數不用看直接選 ReLU 就行了，主要原因還是梯度消失。不過大家可以試試 ReLU 的不同變種，比如

和 MobileNet 用的 ReLU-6 (y = min(max(x, 0), 6))：

訓練過程

▌5.優化器選Adam

這也是我個人的經驗只談。之前用 SGD 經常會卡在局部極小值或者出現梯度爆炸。我推薦大家一開始把學習速率設為 0.001 然后其他參數都用默認：

▌6.動態調整學習速率

一般來說當損失函數不再下降的時候我們就該停止訓練了，因為再訓練就過擬合了。不過如果我們發現損失函數出現上下震蕩的情況，則可能通過減小學習速率讓損失函數變得更小。

下面這個例子中我們可以看到學習速率一開始設的是 0.01，然后從 32 期開始出現震蕩（黃線）。這里通過將學習速率改為 0.001（藍線）使損失函數又減小了大概 0.3。

▌7.權重初始化原則

我個人喜歡把偏置量設為 0，權重則用傳統的正態分布。我一般用的是 Glorot 正態分布初始化法：

▌8.把數據集順序打亂

老生常談了。TensorFlow.js 中我們可以用 tf.utils.shuffle 來實現。

▌9. 保存模型

js 可以通過 FileSaver.js 來實現模型的存儲（或者叫下載）。比如下面的代碼就可以把模型所有的權重保存起來：

保存成什么格式是自己定的，但 FileSaver.js 只管存，所以這里要用JSON.strinfify 把 Blob 轉成字符串：

調試

▌10.保證預處理和后處理的正確性

雖然是句廢話但“垃圾數據垃圾結果”實在是至理名言。標記要標對，每層的輸入輸出也要前后一致。尤其是對圖片做過一些預處理和后處理的話更要仔細，有時候這些小問題還比較難發現。所以雖然費些功夫但磨刀不誤砍柴工。

▌11.自定義損失函數

TensorFlow.js 提供了很多現成的損失函數給大家用，而且一般說來也夠用了，所以我不太建議大家自己寫。如果實在要自己寫的話，請一定注意先測試測試。

▌12.在數據子集試試過擬合

我建議大家模型定義好之后先挑個十幾二十張圖試試看損失函數有沒有收斂。最好能把結果可視化一下，這樣就能很明顯地看出這個模型有沒有成功的潛質。

這樣做我們也能早早地發現模型和預處理時的一些低級錯誤。這其實也就是 11 條里說的測試測試損失函數。

性能

▌13.內存泄漏

不知道大家知不知道 TensorFlow.js 不會自動幫你進行垃圾回收。張量所占的內存必須自己手動調用 tensor.dispose() 來釋放。如果忘記回收的話內存泄漏是早晚的事。

判斷有沒有內存泄漏很容易。大家把 tf.memory() 每次迭代都輸出來看看張量的個數。如果沒有一直增加那說明沒泄漏。

▌14.調整畫布大小，而不是張量大小

在調用 TF . from pixels 之前，要將畫布轉換成張量，請調整畫布的大小，否則你會很快耗盡 GPU 內存。

如果你的訓練圖像大小都一樣，這將不會是一個問題，但是如果你必須明確地調整它們的大小，你可以參考下面的代碼。（注意，以下語句僅在 tfjs - core 的當前狀態下有效，我當前正在使用 tfjs - core 版本 0.12.14）

▌15.慎選批大小

每一批的樣本數選多少，也就是批大小顯然取決于我們用的什么 GPU 和網絡結構，所以大家最好試試不同的批大小看看怎么最快。我一般從 1 開始試，而且有時候我發現增加批大小對訓練的效率也沒啥幫助。

▌16.善用IndexedDB

我們訓練的數據集因為都是圖片所以有時候還是挺大的。如果每次都下載的話肯定效率低，最好是用 IndexedDB 來存儲。IndexedDB 其實就是瀏覽器里嵌入的一個本地數據庫，任何數據都能以鍵值對的形式進行存儲。讀取和保存數據也只要幾行代碼就能搞定。

▌17.異步返回損失函數值

要實時監測損失函數值的話可以用下面的代碼這來自己算然后異步返回：

需要注意的是如果每期訓練完要把損失函數值存到文件里的話這樣的代碼就有點問題了。因為現在損失函數的值是異步返回了所以我們得等最后一個 promise 返回才能存。不過我一般都暴力地在一期結束之后直接等個 10 秒再存：

▌18.權重的量化

為了實現又小又快的目標，在模型訓練完成之后我們應該對權重進行量化來壓縮模型。權重量化不光能減小模型的體積，對提高模型的速度也很有幫助，而且幾乎全是好處沒壞處。這一步就讓模型又能小又能快，非常適合我們在瀏覽器里訓練深度學習模型。

在瀏覽器里訓練深度學習模型的十八招（實際十七招）就總結到這里，希望大家讀了這篇文章能夠有所收獲。

如果有問題也歡迎在后臺給我們留言，大家一起討論！