編者按：opencv4nodejs作者Vincent Mühler分享了基于tensorflow.js設計、訓練面向web的神經網絡模型的經驗。

在將現有的一些目標檢測、臉部檢測、臉部識別的模型移植到tensorflow.js后，我發現有的模型在瀏覽器中的表現沒有達到最優，而另一些模型在瀏覽器上的表現相當不錯。瀏覽器內的機器學習潛力巨大，tensorflow.js等庫給web開發者帶來了太多可能性。

然而，直接在瀏覽器中運行的深度模型也帶來了新的挑戰和限制，畢竟一些現存的模型不是為在客戶端瀏覽器中運行而設計的，更別說移動端瀏覽器了。就拿當前最先進的目標檢測模型來說：它們通常需要大量計算資源才能以合理的fps運行，更別說實時運行了。另外，對一個簡單的web應用來說，分發100MB以上的模型權重至客戶端瀏覽器也不現實。

為web訓練高效的深度學習模型

不過不要放棄希望！基于一些基本原則，我們可以創建和訓練很不錯的模型，為在web環境中運行而優化。信不信由你，實際上我們可以訓練相當不錯的圖像分類模型，甚至是目標檢測模型，大小不過幾兆，甚至幾百K：

這篇文章將給出一些通用的建議，幫助你開始訓練自己的卷積神經網絡（CNN），以及一些基于tensorflow.js為web應用和移動端瀏覽器訓練CNN的建議。

你和我說要在瀏覽器里訓練深度學習模型？

你也許會好奇：為什么要在瀏覽器里基于tensorflow.js訓練我的模型，而不是直接在自己的機器上基于tensorflow訓練模型？你當然可以在自己的機器上訓練，特別是如果你的機器配備了NVIDIA顯卡。tensorflow.js在底層使用了WebGL加速，所以在瀏覽器中訓練模型的一個好處是可以利用AMD顯卡。另外，在瀏覽器中訓練模型，可以更好地保護用戶隱私，更容易讓用戶信任。

網絡架構

顯然，訓練模型之間，首先需要實現模型。通常人們會建議挑選一個久經考驗的現有架構，比如YOLO、SSD、ResNet、MobileNet等。

我個人認為，在設計自己的架構時應用現有架構的一些概念是很有價值的，但直接在web應用中引入這些現存架構可能不妥，因為我們想要尺寸較小、推理較快（理想情況下能實時推理）、容易訓練的模型。

不管你是打算引入一個現有架構，還是完全從頭開始，都可以參考下面的一些建議，在我自己為web設計高效CNN架構時，這些建議幫助很大：

1. 從小型架構開始

牢記，能夠取得較好的精確度的網絡越小，它進行推理的時間越短，用戶下載、緩存模型更容易。此外，較小的模型參數較少，因此在訓練時收斂更快。

如果發現當前的網絡架構表現不是很好，或者沒有達到你想要的精確度，你仍然可以逐漸增加網絡的尺寸，例如，增加每層的卷積過濾器數量，或者直接堆疊更多網絡層，構造更深的網絡。

2. 應用深度可分離卷積

既然我們是要打造一個新模型，毫無疑問我們想要使用深度可分離卷積，而不是傳統的2D卷積。深度可分離卷積將常規的卷積操作分為兩部分，首先分頻道進行卷積，接著應用1x1卷積。和常規卷積操作相比，深度可分離卷積的參數更少，這意味著更少的浮點運算，更容易并行化，推理更快（我曾經碰到過，僅僅將常規卷積替換為深度可分離卷積后，推理就加速了10x），更少的資源占用（在移動設備上意味著性能提升）。此外，因為參數更少，訓練所需時間也更短。

MobileNet和Xception應用了深度可分離卷積（tensorflow.js中的MobileNet和PoseNet）。深度可分離卷積是否導致精確度下降大概沒有定論，但就我的經驗而言，web和移動端模型，毫無疑問應該使用深度可分離卷積。

長話短說：我建議在第一層使用常規的conv2d，畢竟通常而言這一層也沒有多少參數，這樣可以在提取特征中保留RGB通道之間的關系。

export type ConvParams = {

filter: tf.Tensor4D

bias: tf.Tensor1D

}

exportfunction convLayer(

x: tf.Tensor4D,

params: ConvParams,

stride: [number, number],

padding: string

): tf.Tensor4D {

return tf.tidy(() => {

let out = tf.conv2d(x, params.filter, stride, padding)

out = tf.add(out, params.bias)

return out

})

}

剩下的都用深度可分離卷積，用3 × 3 × channelsin × 1過濾器和1 × 1 × channelsin × channels_out過濾器替換單個核。

export type SeparableConvParams = {

depthwise_filter: tf.Tensor4D

pointwise_filter: tf.Tensor4D

bias: tf.Tensor1D

}

exportfunction depthwiseSeparableConv(

x: tf.Tensor4D,

params: SeparableConvParams,

stride: [number, number],

padding: string

): tf.Tensor4D {

return tf.tidy(() => {

let out = tf.separableConv2d(x, params.depthwise_filter: tf.Tensor4D, params.pointwise_filter, stride, padding)

out = tf.add(out, params.bias)

return out

})

}

所以，原本使用tf.conv2d，形狀為[3, 3, 32, 64]的核，現在改用tf.separableConv2d，一個形狀為[3, 3, 32, 1]的核加上一個[1, 1, 32, 64]的核。

3. 跳層連接和密集連接塊

一旦決定創建較深的網絡，很快就會面臨訓練神經網絡最常遇到的問題：梯度消失問題。若干epoch后，損失以極其微小的幅度下降，導致慢得荒謬的訓練，甚至是完全收斂不了。

ResNet和DenseNet使用的跳層連接可以緩解創建更深的架構時遇到的梯度消失問題。我們只需將之前的網絡層的輸出，添加到網絡中較深層的輸入之中（在應用激活函數之前）：

跳層連接背后的直覺是，梯度不必僅僅通過卷積層（或全連接層）反向傳播（這導致梯度漸漸消失）。它們可以添加跳層連接操作，“跳過”一些網絡層。

顯然，跳層連接要工作，需要輸出和輸入之間的形狀互相匹配。假設我們想要跳層連接網絡層A和網絡層B，那么A的輸出形狀需要和B的輸入形狀匹配。如果你想要創建殘差塊或密集連接塊，只需確保卷積過濾器數目一致，步長為1，補齊方案相同。順便說下，還有其他方法，比如補齊A的輸出，使其滿足B的輸入的形狀，或是連接之前層的特征映射，使得相連的層的深度互相匹配。

剛開始，我擺弄了一下類似ResNet的方法，直接隔層引入跳層連接，如上圖所示。不過我很快發現，密集連接塊的效果更好，并且應用密集連接塊后，模型收斂的時間馬上下降了：

這里是一個密集連接塊實現的例子，我在face-api.js的68點臉部識別中將其作為基本構件。這個構件使用了4個深度可分離卷積層（注意，第一個密集塊的第一個卷積層用了常規卷積），每塊的第一個卷積操作步長為2，以縮小輸入：

export type DenseBlock4Params = {

conv0: SeparableConvParams | ConvParams

conv1: SeparableConvParams

conv2: SeparableConvParams

conv3: SeparableConvParams

}

exportfunction denseBlock4(

x: tf.Tensor4D,

denseBlockParams: DenseBlock4Params,

isFirstLayer: boolean = false

): tf.Tensor4D {

return tf.tidy(() => {

const out0 = isFirstLayer

? convLayer(x, denseBlockParams.conv0 as ConvParams, [2, 2], 'same')

: depthwiseSeparableConv(x, denseBlockParams.conv0 as SeparableConvParams, [2, 2], 'same')

as tf.Tensor4D

const in1 = tf.relu(out0) as tf.Tensor4D

const out1 = depthwiseSeparableConv(in1, denseBlockParams.conv1, [1, 1], 'same')

// 第一個連接

const in2 = tf.relu(tf.add(out0, out1)) as tf.Tensor4D

const out2 = depthwiseSeparableConv(in2, denseBlockParams.conv2, [1, 1], 'same')

// 第二個連接

const in3 = tf.relu(tf.add(out0, tf.add(out1, out2))) as tf.Tensor4D

const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv3, [1, 1], 'same')

// 最后一個連接

return tf.relu(tf.add(out0, tf.add(out1, tf.add(out2, out3)))) as tf.Tensor4D

})

}

4. 使用ReLU系列激活函數

除非你有特別的理由，我建議直接使用tf.relu，因為ReLU系列激活函數有助于緩解梯度消失問題。

你也可以試驗ReLU的其他變體，例如leaky ReLU，YOLO架構就用了這個：

exportfunction leakyRelu(x: tf.Tensor, epsilon: number) {

return tf.tidy(() => {

const min = tf.mul(x, tf.scalar(epsilon))

return tf.maximum(x, min)

})

}

或者Mobilenet用的ReLU-6：

exportfunction relu6(x: tf.Tensor) {

return tf.clipByValue(x, 0, 6)

}

訓練

有了初始架構之后，就可以開始訓練模型了。

5. 如果拿不準主意，直接用Adam

剛開始訓練自己的模型的時候，我想知道哪種優化算法是最好的？我從原始的SGD開始，有時候會陷入局部極小值，甚至導致梯度爆炸——模型權重無限增長，逐漸變為NaN.

我并不打算說，對所有問題而言，Adam都是最優選擇，但我發現它是訓練新模型最簡單也最健壯的方式，只需將初始學習率定為0.001，然后以默認參數啟動Adam：

const optimizer = tf.train.adam(0.001)

6. 調整學習率

一旦損失沒有明顯下降，那么我們的模型很可能收斂了（或陷入局部極小值），無法進一步學習了。此時我們也許可以直接停止訓練過程，以免模型過擬合（或者嘗試不同的架構）。

然而，也有可能還有一點壓榨的余地，通過調整（調低）學習率，讓模型進一步學習。特別是在整個訓練集上的總誤差開始震蕩（一會兒高一會兒低）的時候，試著降低學習率也許是個好主意。

下圖是個例子。從第46個epoch開始，損失值開始震蕩，從46個epoch后，將學習率從0.001調低至0.0001，再訓練10個epoch，可以進一步降低總誤差。

7. 權重初始化

如果你對如何恰當地初始化權重毫無頭緒（我剛開始的時候就是這樣），那么可以簡單地將所有偏置初始化為零（tf.zeros），將權重初始化為從某種正態分布中隨機取樣的非零值（比如使用tf.randomNormal）。不過我現在更喜歡用glorot正態分布：

const initializer = tf.initializers.glorotNormal()

const depthwise_filter = initializer.apply([3, 3, 32, 1])

const pointwise_filter = initializer.apply([1, 1, 32, 64])

const bias = tf.zeros([64])

8. 打亂輸入

訓練神經網絡的常見建議是在每個epoch前隨機化訓練樣本的出現順序。我們可以使用tf.utils.shuffle：

exportfunction shuffle(array: any[]|Uint32Array|Int32Array|Float32Array): void

9. 使用FileSaver.js保存模型的快照

FileSaver.js提供了一個saveAs函數，讓我們可以儲存任意類型的文件至下載文件夾。我們可以用它來保存模型權重：

const weights = newFloat32Array([... model weights, flat array])

saveAs(newBlob([weights]), 'checkpoint_epoch1.weights')

我們也可以存儲為json格式，例如保存epoch的累計損失：

const losses = { totalLoss: ... }

saveAs(newBlob([JSON.stringify(losses)]), 'loss_epoch1.json')

調錯

在花費大量時間訓練模型之前，我們想要確保模型確實能夠學習我們預想的東西，并清除任何潛在的錯誤和bug源。下面的幾條建議有助于你避免浪費大量時間訓練出一堆垃圾并懷疑人生：

10. 檢查輸入數據、預處理和后處理邏輯

如果你傳入網絡的是垃圾，網絡返還的也會是垃圾。因此，需要確保輸入數據標注正確，網絡的輸入也符合預期。特別地，如果你實現了隨機剪裁、補齊、截成正方形、居中、減去均值之類的預處理邏輯，別忘了可視化預處理之后的輸入，以檢查輸入是否符合預期。另外，我強烈建議單元測試這些步驟。后處理也是一樣。

我知道這聽起來是一些乏味的額外工作，但它們無疑是有價值的！你不會相信，我花了多少小時嘗試搞明白為什么我的目標檢測器完全學不會檢測面部，直到我逐漸發現，由于錯誤的剪裁和變形，我的預處理邏輯將輸入變成了垃圾。

11. 檢查損失函數

在大多數情況下，tensorflow.js都可以提供所需的損失函數。然而，在你需要實現自己的損失函數的情況下，你絕對需要單元測試！不久前我曾經基于tfjs-core API從頭構建YOLO v2的損失函數，從頭實現損失函數很麻煩，除非你分解問題，并確保每一部分的計算結果符合預期。

12. 先過擬合小數據集

先過擬合訓練數據的一個小子集，從而驗證損失函數可以收斂，模型確實可以學到有用的東西。一般而言，這是一個好主意。例如，你可以直接從訓練數據中選取10到20張圖像，并訓練若干epoch。損失收斂后，在這10到20張圖像上運行推理，并可視化結果：

這是非常重要的一步，有助于消除神經網絡實現中各種造成bug的來源，包括預處理邏輯和后處理邏輯在內。因為，如果代碼中有不少bug，模型不太可能作出符合期望的預測。

特別是實現自己的損失函數的時候，毫無疑問，你希望在正式開始訓練之前確保模型能夠收斂。

性能

最后，我想給出一些有助于降低訓練時間，以及防止瀏覽器因內存泄漏而奔潰的建議。

13. 避免明顯的內存泄露

除非你在tensorflow.js方面完全是新手，你大概已經知道你需要手動丟棄未使用的張量，以釋放它們占用的內存。你可以通過調用tensor.dispose()或將操作封裝在tf.tidy塊中做到這一點。確保不要因為沒有正確地丟棄張量而導致內存泄露，否則你的應用遲早會用盡內存。

識別這類內存泄露很容易。只需在若干次迭代中調用tf.memory()記錄日志，以驗證張量的數量沒有隨著每次迭代而異常增長：

14. 調整canvas的尺寸，而不是張量的尺寸

注意，這條建議只適用于當前的tfjs-core（我使用的版本是0.12.14），因為未來的版本可能會修正這一問題。

我知道這也許聽起來有點奇怪：為什么不使用tf.resizeBilinear、tf.pad等來重整輸入張量的形狀，以匹配網絡輸入要求的形狀？這是因為當前tfjs有一個GPU顯存泄露的bug（#604）

長話短說，在調用tf.fromPixels將canvas轉換成張量前，首先調整canvas的尺寸，以匹配網絡輸入要求的形狀。不這么做的話，你會很快耗盡GPU顯存（取決于訓練數據中圖像尺寸的不同程度）。如果所有訓練圖像尺寸一致，不會有什么問題，但如果尺寸不一致，你可以使用以下代碼調整尺寸：

exportfunction imageToSquare(img: HTMLImageElement | HTMLCanvasElement, inputSize: number): HTMLCanvasElement {

const dims = img instanceof HTMLImageElement

? { width: img.naturalWidth, height: img.naturalHeight }

: img

const scale = inputSize / Math.max(dims.height, dims.width)

const width = scale * dims.width

const height = scale * dims.height

const targetCanvas = document.createElement('canvas')

targetCanvas .width = inputSize

targetCanvas .height = inputSize

targetCanvas.getContext('2d').drawImage(img, 0, 0, width, height)

return targetCanvas

}

15. 找到最優的batch大小

不要把batch大小設得過大！嘗試不同的batch大小，測量反向傳播所需的時間。最優batch大小取決于GPU、輸入尺寸、網絡復雜度。甚至某些情況下，不使用batch，逐一輸入才是最好的。

如果拿不準，我會將batch大小設為1（也就是不使用batch）。我個人發現在一些情形下，增加batch大小對性能提升沒什么幫助，不過，在另一些情形下，我發現，在相當小的網絡上，為112 × 112的圖像創建16到24的batch能帶來大約1.5-2.0的速度提升。

16. 緩存、離線存儲、Indexeddb

訓練圖像（和標簽）也許很多，取決于圖像的尺寸和數量，可能達到1GB甚至更多。由于在瀏覽器下無法直接從磁盤讀取圖像，我們需要使用一個文件代理，比如一個簡單的express服務器，托管訓練數據，然后讓瀏覽器通過文件代理獲取數據。

顯然，這很不高效，不過，別忘了，在瀏覽器中訓練時，如果數據集足夠小，我們大概會把所有數據放在內存中，這顯然也不怎么高效。剛開始，我嘗試增加瀏覽器的緩存大小，以直接將所有數據緩存到磁盤上，不過看起來新版的Chrome和FireFox都不再支持這個功能了。

我最終決定使用Indexddb，這是一個瀏覽器中的數據庫，可以用來存儲整個訓練集和測試集。上手Indexeddb很容易，只需幾行代碼，就可以以鍵值存儲的格式儲存和查詢所有數據。在Indexeddb中，我們可以將標簽存儲為json對象，將圖像數據存儲為blob。

查詢Indexeddb相當快，至少比反復從文件代理查詢并獲取文件要快。另外，將數據轉移到Indexeddb之后，整個訓練過程完全可以離線進行，也就是說，之后的訓練過程不再需要文件代理服務器了。

17. 異步損失報告

這個一個非常有效的簡單技巧，有助于大量減少訓練的迭代次數。如果我們想要獲取optimizer.minimize返回的損失張量的值，以了解訓練過程中的損失值，我們希望避免等待每次迭代的CPU和GPU同步數據。相反，我們希望迭代異步地報告損失：

const loss = optimizer.minimize(() => {

const out = net.predict(someInput)

const loss = tf.losses.meanSquaredError(

groundTruth,

out,

tf.Reduction.MEAN

)

return loss

}, true)

loss.data().then(data => {

const lossValue = data[0]

window.lossValues[epoch] += (window.lossValues[epoch] || 0) + lossValue

loss.dispose()

})

別忘了，現在損失是異步報告的，所以如果我們想要把每個epoch末的總體損失保存到一個文件，我們需要等待最后一個promise得到滿足。我通常直接使用setTimeout在每個epoch完成后的10秒之后再記錄總體損失。

if (epoch !== startEpoch) {

// 丑陋的等待最后一個promise滿足的方法

const previousEpoch = epoch - 1

setTimeout(() => storeLoss(previousEpoch, window.losses[previousEpoch]), 10000)

}

成功訓練模型之后

18. 權重量化

完成模型訓練并對模型的表現滿意后，我建議應用權重量化以壓縮模型大小。通過量化模型權重，可以將模型大小壓縮至1/4！盡可能地壓縮模型大小是很關鍵的，因為它有助于快速地傳輸模型權重至客戶端應用，特別是這樣的壓縮幾乎沒有什么代價。

所以，別忘了參考我寫的tensorflow.js的權重量化指南：https://itnext.io/shrink-your-tensorflow-js-web-model-size-with-weight-quantization-6ddb4fcb6d0d