使用Apache TVM將機器學(xué)習(xí)編譯為WASM和WebGPU
TLDR
在Apache TVM深度學(xué)習(xí)編譯器中引入了對WASM和WebGPU的支持。實驗表明，在將模型部署到Web時，TVM的WebGPU后端可以接近本機 GPU的性能。

概述
計算是現(xiàn)代機器學(xué)習(xí)應(yīng)用程序的支柱之一。GPU的引入加快了深度學(xué)習(xí)的工作量，極大地提高了運行速度。部署機器學(xué)習(xí)的需求不斷增長，瀏覽器已成為部署智能應(yīng)用程序的自然之所。
TensorFlow.js和ONNX.js將機器學(xué)習(xí)引入瀏覽器，但Web版本和本機版本之間在性能上仍然存在著不小的差距。許多原因之一是缺乏對Web上GPU的標(biāo)準(zhǔn)訪問和高性能訪問。WebGL缺少高性能著色學(xué)習(xí)所需的重要功能，例如計算著色器和通用存儲緩沖區(qū)。
WebGPU是下一代Web圖形標(biāo)準(zhǔn)。與最新一代的圖形API（例如Vulkan和Metal）一樣，WebGPU提供了一流的計算著色器支持。
為了探索在瀏覽器中使用WebGPU進行機器學(xué)習(xí)部署的潛力，增強了深度學(xué)習(xí)編譯器Apache（incubating）TVM，以WASM（用于計算啟動參數(shù)并調(diào)用設(shè)備啟動的主機代碼）和WebGPU（用于設(shè)備）為目標(biāo)。執(zhí)行初步結(jié)果是非常積極的-第一次，可以在Web上部署機器學(xué)習(xí)應(yīng)用程序，同時仍能接近GPU的本機性能。
機器學(xué)習(xí)編譯器

試用WebGPU的自然反應(yīng)是為深度神經(jīng)網(wǎng)絡(luò)（矩陣乘法和卷積）中的原始算子編寫著色器，然后直接優(yōu)化性能。這是現(xiàn)有框架（例如TensorFlow.js）使用的傳統(tǒng)工作流程。
相反，采用了基于編譯的方法。TVM自動從TensorFlow，Keras，PyTorch，MXNet和ONNX等高級框架中提取模型，使用機器學(xué)習(xí)驅(qū)動的方法自動生成低級代碼，在這種情況下，將以SPIR-V格式計算著色器。然后可以為可部署模塊生成的代碼打包。
編譯的方法的一個重要優(yōu)點是基礎(chǔ)架構(gòu)的重用。通過重用基礎(chǔ)結(jié)構(gòu)來優(yōu)化CUDA，Metal和OpenCL等本機平臺的GPU內(nèi)核，能夠輕松地（相對于其它方法）以Web為目標(biāo)。如果WebGPU API到本機API的映射有效，可以通過很少的工作獲得類似的性能。更重要的是，AutoTVM基礎(chǔ)架構(gòu)，能夠針對特定模型專門化計算著色器，從而能夠為感興趣的特定模型生成最佳的計算著色器。
構(gòu)建WASM和WebGPU編譯器
為了構(gòu)建可以針對WASM和WebGPU的編譯器，需要以下元素：
? 用于計算著色器的SPIR-V生成器。
? 主機程序的WASM生成器。
? 加載和執(zhí)行生成的程序的runtime。
TVM已經(jīng)有Vulkan的SPIR-V目標(biāo)，使用LLVM生成主機代碼。可以僅將二者的用途重新生成設(shè)備和主機程序。
主要挑戰(zhàn)是runtime。需要一個runtime來加載著色器代碼，并使主機代碼對話能夠正確地與著色器通信。TVM具有最低的基于C ++的runtime。構(gòu)建了一個最小的Web runtime庫，生成的著色器和主機驅(qū)動代碼鏈接，生成一個WASM文件。但是，此WASM模塊仍然包含兩個未知的依賴項：
? runtime需要調(diào)用系統(tǒng)庫調(diào)用（malloc，stderr）。
? wasmruntime需要與WebGPU驅(qū)動程序進行交互（在Javascript中，WebGPU API是the first-class citizen）。
WASI是解決第一個問題的標(biāo)準(zhǔn)解決方案。盡管網(wǎng)絡(luò)上還沒有成熟的WASI，使用emscripten生成類似WASI的庫，提供這些系統(tǒng)庫。
通過在TVM的JS runtime內(nèi)部構(gòu)建WebGPU runtime來解決第二個問題，在調(diào)用GPU代碼時，從WASM模塊中回調(diào)這些功能。使用TVM runtime系統(tǒng)中的PackedFunc機制，可以通過將JavaScript閉包傳遞到WASM接口，直接公開高級runtime原語。這種方法將大多數(shù)runtime代碼保留在JavaScript中，隨著WASI和WASM支持的成熟，可以將更多JS代碼引入WASM runtime。

性能

進行了一個快速實驗，比較了通過TVM的WebGPU后端和使用本地GPU runtime（Metal和OpenCL）的本地目標(biāo)執(zhí)行完整計算圖的情況。在MobileNet模型上，可以發(fā)現(xiàn)WebGPU可以接近Metal的性能。假設(shè)Chrome WebGPU的runtime以MacOS上的Metal（而不是OpenCL）為目標(biāo)，可以放心地假設(shè)以GPU為目標(biāo)時，性能幾乎沒有損失。
此基準(zhǔn)不包括CPU到GPU的數(shù)據(jù)復(fù)制成本，而僅基準(zhǔn)GPU的執(zhí)行。從CPU到GPU的數(shù)據(jù)復(fù)制，仍會占用25％的執(zhí)行時間。可以通過諸如連續(xù)執(zhí)行設(shè)置中的雙緩沖之類的方法，進一步攤銷這些成本。
報告的mobilenet的端到端runtime，絕不是最佳選擇，重復(fù)使用了GTX 1080 Ti的優(yōu)化程序，這與Intel圖形GPU截然不同。希望通過在目標(biāo)平臺上使用AutoTVM來進一步提高性能。
展望未來
結(jié)果表明，在網(wǎng)絡(luò)上進行機器學(xué)習(xí)有許多有趣的機會。值得注意的是，WebGPU是一個仍在不斷發(fā)展的API，其含義可能會超出Web應(yīng)用程序。例如，當(dāng)WebGPU成熟，通過WASI標(biāo)準(zhǔn)化時，可以將其定位為WebGPU的本機API，使用WebGPU的獨立WASM應(yīng)用程序。
TVM社區(qū)還積極地在基于Rust的runtime上工作，該runtime將提供更強大的WASM支持，wgpu和Rust WASM生態(tài)系統(tǒng)等項目的交互更加輕松。
提出的方法為大多數(shù)WASM的應(yīng)用場景提供了有效的機器學(xué)習(xí)支持。接近本機的性能，可以釋放瀏覽器上更好的聯(lián)合學(xué)習(xí)功能。相同的編譯程序包，也應(yīng)該能夠在本機WASM執(zhí)行程序上運行，為應(yīng)用程序提供sandbox 。

總結(jié)

以上是生活随笔為你收集整理的使用Apache TVM将机器学习编译为WASM和WebGPU的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。