TensorFlow.js 可在瀏覽器和 Node.js 中運作,而且這兩個平台都有許多不同的可用組態。每個平台都有一組獨特的考量因素,會影響應用程式的開發方式。
在瀏覽器中,TensorFlow.js 支援行動裝置和桌上型電腦。每個裝置都有一組特定的限制,例如可用的 WebGL API,系統會自動判斷並為您設定這些限制。
在 Node.js 中,TensorFlow.js 支援直接繫結至 TensorFlow API,或使用速度較慢的原始 CPU 實作。
環境
當執行 TensorFlow.js 程式時,特定的組態稱為環境。環境包含單一的全域後端,以及一組控制 TensorFlow.js 細微功能的旗標。
後端
TensorFlow.js 支援多個不同的後端,這些後端實作張量儲存和數學運算。在任何給定時間,只有一個後端處於作用中狀態。在大多數情況下,TensorFlow.js 會根據目前的環境自動為您選擇最佳後端。然而,有時瞭解正在使用的後端以及如何切換後端非常重要。
若要尋找您正在使用的後端
console.log(tf.getBackend());
如果您想要手動變更後端
tf.setBackend('cpu');
console.log(tf.getBackend());
WebGL 後端
WebGL 後端「webgl」目前是瀏覽器最強大的後端。此後端的速度比原始 CPU 後端快達 100 倍。張量會儲存為 WebGL 紋理,而數學運算則在 WebGL shader 中實作。以下是使用此後端時需要知道的一些實用資訊:\
避免封鎖 UI 執行緒
當呼叫運算 (例如 tf.matMul(a, b)
) 時,會同步傳回產生的 tf.Tensor
,但矩陣乘法的計算可能尚未實際完成。這表示傳回的 tf.Tensor
只是計算的控制代碼。當您呼叫 x.data()
或 x.array()
時,值會在計算實際完成時解析。這使得使用非同步 x.data()
和 x.array()
方法 (而非同步對應方法 x.dataSync()
和 x.arraySync()
) 變得非常重要,以避免在計算完成時封鎖 UI 執行緒。
記憶體管理
使用 WebGL 後端時需要注意的一點是,需要明確的記憶體管理。WebGLTexture (張量資料最終儲存的位置) 不會由瀏覽器自動進行垃圾收集。
若要銷毀 tf.Tensor
的記憶體,您可以使用 dispose()
方法
const a = tf.tensor([[1, 2], [3, 4]]);
a.dispose();
在應用程式中將多個運算鏈結在一起非常常見。保留對所有中間變數的參照以處置它們可能會降低程式碼可讀性。為了解決這個問題,TensorFlow.js 提供了 tf.tidy()
方法,該方法會清除函數執行後未傳回的所有 tf.Tensor
,這類似於在執行函數時清除區域變數的方式
const a = tf.tensor([[1, 2], [3, 4]]);
const y = tf.tidy(() => {
const result = a.square().log().neg();
return result;
});
精確度
在行動裝置上,WebGL 可能僅支援 16 位元浮點紋理。然而,大多數機器學習模型都使用 32 位元浮點權重和啟動進行訓練。當將模型移植到行動裝置時,這可能會導致精確度問題,因為 16 位元浮點數只能表示範圍 [0.000000059605, 65504]
內的數字。這表示您應注意模型中的權重和啟動不要超出此範圍。若要檢查裝置是否支援 32 位元紋理,請檢查 tf.ENV.getBool('WEBGL_RENDER_FLOAT32_CAPABLE')
的值,如果為 false,則表示裝置僅支援 16 位元浮點紋理。您可以使用 tf.ENV.getBool('WEBGL_RENDER_FLOAT32_ENABLED')
來檢查 TensorFlow.js 目前是否正在使用 32 位元紋理。
Shader 編譯和紋理上傳
TensorFlow.js 透過執行 WebGL shader 程式,在 GPU 上執行運算。這些 shader 會在使用者要求執行運算時延遲組裝和編譯。shader 的編譯發生在主要執行緒的 CPU 上,速度可能會很慢。TensorFlow.js 會自動快取編譯後的 shader,使第二次以相同形狀的輸入和輸出張量呼叫相同運算的速度快得多。通常,TensorFlow.js 應用程式會在應用程式的生命週期內多次使用相同的運算,因此第二次通過機器學習模型的速度會快得多。
TensorFlow.js 也將 tf.Tensor 資料儲存為 WebGLTexture。當建立 tf.Tensor
時,我們不會立即將資料上傳到 GPU,而是將資料保留在 CPU 上,直到 tf.Tensor
在運算中使用為止。如果 tf.Tensor
第二次使用,則資料已在 GPU 上,因此沒有上傳成本。在典型的機器學習模型中,這表示權重會在第一次通過模型進行預測期間上傳,而第二次通過模型的速度會快得多。
如果您關心第一次通過模型或 TensorFlow.js 程式碼進行預測的效能,我們建議在使用真實資料之前,先傳遞相同形狀的輸入張量來預熱模型。
例如
const model = await tf.loadLayersModel(modelUrl);
// Warmup the model before using real data.
const warmupResult = model.predict(tf.zeros(inputShape));
warmupResult.dataSync();
warmupResult.dispose();
// The second predict() will be much faster
const result = model.predict(userData);
Node.js TensorFlow 後端
在 TensorFlow Node.js 後端「node」中,TensorFlow C API 用於加速運算。如果可用,這將使用機器的可用硬體加速 (例如 CUDA)。
在此後端中,就像 WebGL 後端一樣,運算會同步傳回 tf.Tensor
。然而,與 WebGL 後端不同的是,運算會在您取回張量之前完成。這表示呼叫 tf.matMul(a, b)
將會封鎖 UI 執行緒。
因此,如果您打算在生產環境應用程式中使用此後端,則應在工作執行緒中執行 TensorFlow.js,以免封鎖主要執行緒。
如需 Node.js 的詳細資訊,請參閱本指南。
WASM 後端
TensorFlow.js 提供 WebAssembly 後端 (wasm
),其提供 CPU 加速,並且可以用作原始 JavaScript CPU (cpu
) 和 WebGL 加速 (webgl
) 後端的替代方案。若要使用它
// Set the backend to WASM and wait for the module to be ready.
tf.setBackend('wasm');
tf.ready().then(() => {...});
如果您的伺服器在不同的路徑或不同的名稱上提供 .wasm
檔案,請在使用後端初始化之前使用 setWasmPath
。如需更多資訊,請參閱 README 中的「使用 Bundler」章節
import {setWasmPath} from '@tensorflow/tfjs-backend-wasm';
setWasmPath(yourCustomPath);
tf.setBackend('wasm');
tf.ready().then(() => {...});
為何選擇 WASM?
WASM 於 2015 年推出,作為一種新的網路型二進位格式,為以 JavaScript、C、C++ 等語言編寫的程式提供在網路上執行的編譯目標。自 2017 年以來,Chrome、Safari、Firefox 和 Edge 支援 WASM,並且全球 90% 的裝置都支援 WASM。
效能
WASM 後端利用 XNNPACK 程式庫 來最佳化神經網路運算子的實作。
與 JavaScript 相比:對於瀏覽器而言,WASM 二進位檔案通常比 JavaScript 套件更快地載入、剖析和執行。JavaScript 是動態類型且具有垃圾收集功能,這可能會導致執行階段速度變慢。
與 WebGL 相比:對於大多數模型而言,WebGL 比 WASM 快,但對於微型模型而言,由於執行 WebGL shader 的固定額外負荷成本,WASM 可能會勝過 WebGL。「何時應使用 WASM」章節將討論做出此決策的啟發法。
可攜性和穩定性
WASM 具有可攜式 32 位元浮點算術,可在所有裝置上提供精確度同位。另一方面,WebGL 是硬體特定的,不同的裝置可能具有不同的精確度 (例如,在 iOS 裝置上回退到 16 位元浮點數)。
與 WebGL 一樣,WASM 受到所有主要瀏覽器的正式支援。與 WebGL 不同的是,WASM 可以在 Node.js 中執行,並且可以在伺服器端使用,而無需編譯任何原生程式庫。
何時應使用 WASM?
模型大小和運算需求
一般而言,當模型較小或您關心缺少 WebGL 支援 (OES_texture_float
擴充功能) 或 GPU 功能較弱的低階裝置時,WASM 是不錯的選擇。下圖顯示了在 2018 MacBook Pro 上 Chrome 中 5 個我們正式支援的 模型在 WebGL、WASM 和 CPU 後端之間的推論時間 (截至 TensorFlow.js 1.5.2)
較小的模型
模型 | WebGL | WASM | CPU | 記憶體 |
---|---|---|---|---|
BlazeFace | 22.5 毫秒 | 15.6 毫秒 | 315.2 毫秒 | .4 MB |
FaceMesh | 19.3 毫秒 | 19.2 毫秒 | 335 毫秒 | 2.8 MB |
較大的模型
模型 | WebGL | WASM | CPU | 記憶體 |
---|---|---|---|---|
PoseNet | 42.5 毫秒 | 173.9 毫秒 | 1514.7 毫秒 | 4.5 MB |
BodyPix | 77 毫秒 | 188.4 毫秒 | 2683 毫秒 | 4.6 MB |
MobileNet v2 | 37 毫秒 | 94 毫秒 | 923.6 毫秒 | 13 MB |
上表顯示,WASM 在各個模型上的速度比純 JS CPU 後端快 10-30 倍,並且在較小的模型 (例如 BlazeFace) 上與 WebGL 競爭,BlazeFace 體積輕巧 (400KB),但具有相當多的運算 (~140)。鑑於 WebGL 程式的每個運算執行都有固定的額外負荷成本,這解釋了為何 BlazeFace 等模型在 WASM 上更快。
這些結果會因您的裝置而異。判斷 WASM 是否適合您的應用程式的最佳方法是在我們的不同後端上進行測試。
推論與訓練
為了處理預先訓練模型部署的主要使用案例,WASM 後端開發將優先考慮推論而非訓練支援。請參閱 WASM 中支援的運算子的最新清單,如果您的模型有不受支援的運算子,請告訴我們。對於訓練模型,我們建議使用 Node (TensorFlow C++) 後端或 WebGL 後端。
CPU 後端
CPU 後端「cpu」是效能最低的後端,但它是最簡單的後端。所有運算都在原始 JavaScript 中實作,這使得它們較難以平行化。它們也會封鎖 UI 執行緒。
此後端對於測試或在 WebGL 無法使用的裝置上非常有用。
旗標
TensorFlow.js 有一組環境旗標,這些旗標會自動評估並判斷目前平台中的最佳組態。這些旗標大多是內部的,但一些全域旗標可以使用公用 API 控制。
tf.enableProdMode():
啟用生產環境模式,這將移除模型驗證、NaN 檢查和其他正確性檢查,以提高效能。tf.enableDebugMode()
:啟用偵錯模式,這會將執行的每個運算以及執行階段效能資訊 (例如記憶體用量和核心總執行時間) 記錄到主控台。請注意,這會大大降低應用程式的速度,請勿在生產環境中使用。