Tensorflow JS 初探 — YOLO

Published in

Cubo AI

6 min readMay 19, 2018

今年 3 月底時，Google 發表了 Tensorflow Javascript 版本 (其實是直接併了人家 deeplearn.js)，提供與 C++ / Python 一致的 API 介面，讓開發者也可以在使用者的瀏覽器上訓練 AI 模型或者提供 inference (好啦，其實我搞不懂有什麼 use case 是要在別人家的瀏覽器上訓練模型)

不過用瀏覽器來載入 model ，並且提供 inference 還蠻有趣的，有些好處，像是減輕一些負載到使用者端，或者如果有隱私上的考量，可以避免傳遞使用者資料到後端。

Tensorflow JS 官網有提供一些範例，不過這裡我想要來試試我比較熟悉的 YOLO ，關於 YOLO 細節可以參考我前篇文章以及這篇。

你可以到這裡來玩看看 👉 https://syshen.github.io/tfjs-playground/

實作是參考 tfjs-yolo-tiny npm 而來。

此外，底下我會改用 tfjs 來簡稱 Tensorflow JS 。

基本介紹

Tensor 即是張量，簡單來說就是矩陣，不管是一維矩陣，還是多維，在深度學習裡，絕大多數都是在做矩陣運算，所以 tensorflow 提供了許多矩陣運算元。在 tfjs 中提供五種 Tensor，分別是 tf.Scalar、tf.Tensor1d、tf.Tensor2d、tf.Tensor3d 跟 tf.Tensor4d，scalar就是純量，其他分別代表1到4維陣列。

tfjs 底層透過 WebGL 存取 GPU 資源，所以運算速度才可以快，所以如果要將優化 tfjs 的程式，最好轉換成 tf.Tensor 來作運算，所以要熟矩陣運算。

為避免太多 tensor 運算，吃掉過多記憶體，tfjs 提供兩個函式讓你可以 free 掉不用的 tensor：

dispose: 呼叫這個函式會立即將該 tensor 從 GPU 記憶體中移除
tf.tidy( fn ): 如果你要執行數個運算，可以透過 tf.tidy ，將你要執行的運算包在傳遞給他的 function 當中。tidy 會在執行完畢時，將函式執行過程中建立的 tensor 從記憶體中移除，除了你回傳的 tensor 之外。

所以，為避免使用者的瀏覽器在執行的 model 過程中爛掉，記得靈活使用這兩個函式來管理你的記憶體使用量。

載入預先訓練好的模型

你不能直接用 Tensorflow 的模型檔，必須經過 tfjs-converter 轉換才可用在 tfjs 的程式中。 tfjs-converter 支援 Tensorflow 的 SavedModel pb 檔，也支持 Keras 的 hdf5 檔(好消息)，但由於 tfjs 支持的運算元有限，不是每個模型檔都能夠順練轉換的，幸好 YOLOv2 不在此列。

YOLO原本是用 darknet 實作的，所以需先轉成 Keras 的 hdf5 檔，再轉成 tfjs 格式。

首先，先下載 YOLOv2 Tiny 設定檔與 weights 檔

下載 yad2k，並轉換成 Keras 檔。不過在轉換前，請先讀一下這個 Open Issue 的討論，yad2k 在轉換 tiny YOLO 的 weights 檔有 bug，請 follow 這個 Open Issue 的討論，先改過 yad2k.py 之後，再執行底下命令 (這個 bug 卡了我好久 XD)

安裝 tensorflowjs converter，並將 Keras H5 檔轉換成 tfjs 格式

接著找個網站 host 你的 model file 就好，我是直接用 github 來 host 這些 model file。

YOLO 辨識與後處理

在程式中載入模型後，便可以透過 predict 函式去偵測，但是在這之前，你必須先將 ImageData 轉成 Tensor 作為輸入，而 YOLO 的輸入張量是 416 x 416 x 3，也就是長寬各 416，RGB 3個 channel 。

但是由於圖檔不一定是方形的，所以我們必然要對圖檔作中間部位的裁切成 416 x 416 大小，然後再做正規化。

predict 出來的結果，會包含 13 x 13 個 grid，每個 grid 各自包含一個矩陣，裡面數據依序是： x, y, w, h, confidence，然後是 coco 80 個類別的 confidence，所以總共是 85 個數據，每個 grid 會有 5 種 anchor box，所以 85 *5 = 425 。

所以 predict 結果的 shape 是 1 x 13 x 13 x 425

後處理就是要將這 13 x 13 x 5 個 box 的 confidence 跟 bounding box 取出來，只留下 confidence 夠高的 bounding box，然後對剩下的 bounding box 作合併，得到的結果就會是我們想要的。

之前在上一篇 YOLO with CoreML 中，直接用 for 迴圈跑完 13 x 13 x 5 次，所以效率很差。如果用這個方法，在網頁上測試結果，predict 花不到一秒鐘，但是計算 bounding box 卻是 predict 的 10 幾倍 (暈倒~

就像前面提到的，最好的作法是直接用 Tensor 運算，可以運用到 GPU 資源，會快很多。

我知道這段 code 不好懂，但是解釋起來也麻煩，反正最後會得到數個候選框(box)，各個候選框的類別 index 、score 、以 x,y 位置與長寬。

不過這些候選框有些有高度重疊，所以還需要經過 Non-max Suppression 來濾掉一些高度重疊的候選框：

只要是 IOU 重疊有超過 0.5 以上的我們就把他過濾掉。

幾個測試結果: