Inference on the Browser with Tensorflow JS

以 CNN 貓狗分類器為例，讓 AI 以不同角度切入前端戰場

11 min readSep 10, 2019

--

全民人工智慧時代來臨，近期正進行公司的深度學習課程中，所以挑了一組常見的 Kaggle Dogs vs. Cats Dataset 來練習。網路上的學習資源已經相當豐富，這篇文章主要會紀錄從訓練端到終端上的過程，以 Python Karas 做訓練，輸出轉換成 Tensorflow.js 的格式，在瀏覽器端進行推論，完整地走過一遍。

Links: GitHub | Jupyter Notebook | Demo

Convolutional Neural Network

針對影像問題上主要會拿 Convolutional Neural Network 來分析，CNN 使用 Convolutional 概念擷取特徵、Pooling 降維加速計算，並在訓練中自動學習最佳參數。而其中如何去設計中間層是非常需要經驗的，雖然可以從頭搭建出自己的模型，但也有更有效率的做法，也就是 Transfer Learning：直接採用 Pre-trained 的模型與參數。本篇文章挑選了由 Keras 實作 ImageNet 資料集所訓練出來的 VGG16 模型，直接複用中間層來萃取圖像特徵，最後再客製化輸出層來作為最後分類器使用。

Training

Running Jupyter Notebook on Google Colab

Jupyter Notebook 是很常見來撰寫程式的工具，所見即所得的特性，相當適合拿來對資料結構做前處理。Notebook 也便於結果的分享，因此 Kaggle 會有許多熱心的人上傳各自的實驗過程，不管是要 Reproduce 或在學習上都很有幫助，筆者也是直接參考推薦數最多的 Uysim’s Notebook 來整理資料。以下是最後整份 Notebook 的紀錄，可以直接跑在 Google Colab 並且選擇 GPU 來加速運算，接下來的段落會把每個步驟稍微解釋：

Jupyter Notebook: https://github.com/evenchange4/nextjs-tfjs-cnn/blob/master/python/vgg16.ipynb

Pre-requirements

在 Notebook 中可以直接下 kaggle CLI 指令，將資料集下載到 Google Colab 的空間，但每次重新連接可能會被分配新的環境，過程中資料不會被保留，所以也可以透過 Mount Google drive 把最後要輸出的結果放到個人的空間上：

!kaggle competitions download -c dogs-vs-cats
drive.mount('/content/drive')

Build Model

Transfer Learning Model 實作上直接拿 Keras 提供的 VGG16，因為設定 include_top=False 忽略 Output Layers，只拿中間層來當作 Feature Extractor，並且設定 trainable = False Freeze Layers，直接使用訓練好的 weights=’imagenet’ 參數。後面加上 Classifier Layers，這幾層的參數則是需要在過程中被從頭訓練的，其中使用 BatchNormalization 讓數值不會過大過小，加速收斂。接著數層全連接層 Dense，以及避免 Overfitting 使用 Dropout：一定機率讓神經元被忽略。最後我們要分類貓或狗，設定輸出兩個維度，並挑選 Softmax 激活函數使得機率總和為 1：

pre_trained_model = VGG16(include_top=False, weights='imagenet')
for layer in pre_trained_model.layers:
    layer.trainable = Falsex = BatchNormalization()(pre_trained_model.output)
x = Flatten()(x)
x = Dense(units=2048, activation="relu")(x)
x = Dense(units=256, activation="relu")(x)
x = Dropout(0.5)(x)
x = Dense(units=2, activation="softmax")(x)model = Model(inputs=pre_trained_model.input, outputs=x)
model.compile(loss='categorical_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])

Prepare data

在真正開始跑訓練前，需要把資料做好準備，首先將下載的資料集切分為訓練用的 Training data，與小部分判斷訓練結果的 Validation data。若資料量不足，為了避免 Overfitting 的狀況，可以使用由 Keras 提供的 Data augmentation ImageDataGenerator 功能，在每一輪 Epoch 自動生成相仿的圖片供訓練。見下圖範例，將原始圖像旋轉、縮放、裁減、翻轉，轉化後的假資料，以提升模型的泛化能力：

Fit Model

跑訓練與調校模型參數是過程中最耗時的，可以採用 GPU 來加速運算，其中設定 batch_size 讓 Mini-batch Gradient Descent 小批小批更新學習參數；透過 epochs 設定資料集要拿來訓練幾輪。除此之外，可以預先設定訓練結束的 callbacks 邏輯，例如若連續幾輪都沒獲得優化，EarlyStopping 就直接中斷訓練；又或是下左圖觀測 Validation Accuracy 連續兩輪都沒有提升，讓 ReduceLROnPlateau 調降 Learning Rate，來減小訓練步伐。下右圖為視覺化最後訓練出來 Accuracy 的變化：

EarlyStopping(patience=5)
ReduceLROnPlateau(monitor='val_acc', 
                  patience=2, 
                  verbose=1, 
                  factor=0.5, 
                  min_lr=0.00001)

Fit model with callbacks & Virtualize training accuracy

Predict Result

最後我們來看沒有提供答案的 Testing data，將預測得到的結果整理好 CSV 格式，上傳到 Kaggle Late Submission 進行 Log Loss 評分，如果滿意可以 model.save 將模型及訓練參數輸出，最後會得到一份 HDF5 格式的檔案 model.h5。

Client side

緊接著進入前端 JavaScript 的世界，使用 Next.js 快速搭建環境，並使用 React Material UI 做了基本 Layout。載入模型前，首先將上個段落的 model.h5 透過 tensorflowjs_converter 轉成 Tensorflow.js 的格式，其輸出包含 model.json 與 Binary weight files，這邊拿準備好的 Docker image 環境來執行 CLI，最後下圖利用 Netron 來視覺化模型，驗證轉換後的結果是否正確：

$ docker run -it --rm \
  -v "$PWD/python:/python" \
  evenchange4/docker-tfjs-converter \
  tensorflowjs_converter --input_format=keras model.h5 model-tfjs

在跑預測前，需要把輸入圖檔調整為模型的 Input Tensor，並注意也需要用當初在 Keras VGG16 的 Image Preprocess 步驟，將圖片 Tensor 做一些前處理運算，因為 Tensorflow.js 的支援，讓這些步驟簡單很多：

// preprocessImage.tsconst imageTensor = tf.browser
  .fromPixels(image)
  .resizeNearestNeighbor([224, 224])
  .toFloat()
  .sub(meanImageNetRGB)
  .reverse(2)
  .expandDims();

Web Worker

最後實驗嘗試將 Predict 函式丟到 Web Worker，使用 WebGL 運算，目前 Chrome 支援度會比較好，但在 CNN 的支援上不是很到位，UI Main thread 還是會有卡頓感，這也是之後前端可以優化的目標。下圖是操作流程，先透過 Web Worker 將模型與參數下載，接著就可以上傳圖片進行瀏覽器端的預測了：

// inference.workerconst model = await tf.loadLayersModel('/static/model.json');
model.predict(imageTensor);

Conclusion

在這篇文章我們走過了 Transfer Learning，一開始複用 VGG16 來搭建一個分類器，使用 Kaggle Dataset 在 Google colab 上用 GPU 來快速訓練，最後把輸出的模行轉換成 Tensorflow.js 的格式，在瀏覽器做本地端預測。雖然目前結果不甚好，之後需要花費大量時間調校，但在深度學習的研究成果下，類神經網路已經非常便於使用了。另外可以研究的方向是，預測準確度與模型大小的取捨，目前 VGG16 輸出檔案太大，造成瀏覽器端需要花費很多時間下載。期待接下來也許可以透過這樣的架構，從生活中或公司的資料實際找尋更有趣的應用案例，如果你喜歡這類的文章，可以 👏 讓我知道，也歡迎訂閱 Medium 喔！

References

Special thanks to mnicnc404.