Go 機器學習框架之火重燃，Google 前研究員開源期望媲美 Jax 的 GoMLX

Go 作爲一門兼具高性能與簡潔性的編程語言，近年來在各種領域得到廣泛應用。然而，在機器學習領域，Go 相比 Python、C++、Julia 等語言，生態仍然較爲薄弱。目前的 Go 機器學習框架無論在功能全面性上，還是在社區生態支持上都難以與 TensorFlow、PyTorch、Jax 等重量級框架抗衡。究其原因，筆者覺得還是 Go 社區缺少熟悉和精通機器學習方面的人才。

不過，隨着時間的推移，總會有一些開發者會嘗試探索在 Go 語言中實現高效機器學習的可能性。近期，筆者就發現一個由 Google 前研究員 Jan Pfeifer[1] 開源的 Go 機器學習框架項目 GoMLX[2]，雖然我不是很瞭解機器學習這個領域，但也能看出一些該項目與以往 Go 機器學習項目的不同。

下面我就來簡單介紹一下 gomlx 項目。

注：該項目開源時間不長，正處於初期開發階段，還是一個實驗性項目，也基本是由 Jan Pfeifer 一人完成。

一. GoMLX 簡介

GoMLX 是一個主要由 Go 實現的、快速且易用的機器學習和通用數學庫與工具集，你可以將其視爲 Go 語言版本的 PyTorch/Jax/TensorFlow(也許現在還不是，但其目標定位是這樣的)。它力求成爲 Go 語言中的標準化機器學習框架，讓開發者在使用 Go 進行機器學習時，無需依賴其他語言或工具鏈。

結合該項目的 README 和作者提供的 GoMLX 教程 [3]，我們可以總結一下目前 GoMLX 目前已具備的主要功能特性，以及其設計上的優點。

GoMLX 已具備的主要功能特性包括：

• 計算圖構建：支持創建和操作計算圖。

• JIT 編譯：能夠即時編譯計算圖以提高執行速度。

• 多後端支持：可以使用不同的後端執行操作，默認使用 XLA。

• 張量操作：提供了基本的張量操作和管理功能。

• 自動微分：支持靜態梯度計算。

• 變量管理：通過 Context 對象管理變量和模型參數。

• 機器學習層：提供了多種常用的機器學習層，如 Dense、Dropout、BatchNorm 等。

• 訓練工具：包括 Trainer 和 Loop 等用於模型訓練的工具。

• 數據集接口：定義了 Dataset 接口用於數據輸入。

• 優化器、損失函數和指標：提供了常用的優化器、損失函數和評估指標。

• 可視化支持：能夠繪製訓練過程中的指標圖表。

• 調試工具：提供了多種調試方法，如形狀斷言、錯誤處理等。

• 預訓練模型支持：提供了諸如 InceptionV3 等預訓練模型。

GoMLX 在設計上相對於其他 Go 機器學習框架的優點有以下幾點：

• 模塊化設計：將不同功能封裝在獨立的包中，便於理解和使用。

• 性能：通過使用 JIT 編譯和 XLA 後端，實現了高性能計算。

• 錯誤處理：採用異常機制處理錯誤，提供了詳細的堆棧跟蹤，有助於調試。

• 文檔完善：強調了文檔的重要性，提供了詳細的 API 文檔和教程。

• 上下文管理：通過 Context 對象統一管理變量、參數和作用域，簡化了模型開發。

• 硬件適配：支持 CPU 和 GPU，並可能擴展到其他硬件（如 TPU）。

• 開發體驗：提供了交互式開發環境（如 Jupyter notebook 支持，通過 GoNB 項目 [4]，作者的另外一個開源項目），方便實驗和學習。

• 可視化集成：內置了訓練過程的可視化工具，方便監控和分析。

二. GoMLX 的長期目標

作者爲 GoMLX 項目設定了明確的長期目標，包括：

• 在 Go 中構建和訓練模型，注重簡單易讀、API 清晰、可組合性強。

• 成爲高效的研究和教育平臺，方便實驗新的 ML 想法

• 成爲可靠的生產平臺，支持：

• 現代加速硬件如 TPU 和 GPU

• 與 TensorFlow Serving 等工業工具集成

• 從 Hugging Face Hub 等導入預訓練模型

• 將模型編譯爲二進制庫或 WebAssembly

三. GoMLX 的 TODO list

前面說過，該項目處於初期開發階段，還有很多待開發的功能，作者在項目中放置了一個 todo list，我們也可以看一下作者對 GoMLX 的功能規劃：

1. 模型相關

• 增加 CNN 相關功能，如 SeparableConv2D，以支持完整的 ResNet50 實現

• 添加 GNN(圖神經網絡) 層

• 支持從 TensorFlow/Jax/PyTorch/Hugging Face 導入模型：

• 支持黑盒推理和進一步微調

• 提供一些標準模型如 ResNet50、ViT、BERT、Chinchilla 等

• 重點支持從 Hugging Face 導入模型

• 擴展計算圖操作和操作工具

• 實現 NaN 和 Inf 值的首次出現檢測，並提供堆棧跟蹤

3. 基礎設施

• 模型保存 / 加載：

• 支持導出爲 TensorFlow 的 "SavedModel" 格式

• 開發僅用於推理的輕量版本

• 支持將計算圖預先編譯 (AOT) 爲不依賴 XLA 的庫

5. 分佈式訓練

• 實現同步鏡像策略 (數據並行)

• 支持多設備 (GPU) 設置

• 支持多租戶 (多主機) 設置

7. 底層支持

• 完善多設備 (如多 GPU) 支持

• 實現自動分佈在多設備上的 DistributedTrainer

• 新後端支持：

• 開發兼容 WebAssembly 的後端 (如 WebGL)，以支持瀏覽器中運行和訓練

• 開發更快 / 更輕的 CPU 後端 (可能使用 ggml)

9. API 改進

• 雖然沒有具體列出，但文檔提到 API 仍有改進空間，預計在早期階段會有一些變化

這些計劃涵蓋了從模型功能擴展到基礎設施改進，以及分佈式和新平臺支持等多個方面，旨在使 gomlx 成爲一個更全面、強大的機器學習框架。

四. 小結

通過上面的介紹，我們看到：GoMLX 站在了 TensorFlow/PyTorch/Jax 等機器學習框架巨頭的肩膀之上，借鑑了他們的設計與優勢，並充分複用已有的工具和庫，結合 Go 語言的特點，但又不侷限在 Go 之中。

GoMLX 的作者就是機器學習研究員，十分清楚機器學習框架用戶的訴求，提供了 GoNB(A Modern Go Kernel for Jupyter Notebooks)，讓用戶可以在熟悉的 Jupyter Notebooks 中完成相關實驗，並內置了圖形化展示工具。GoMLX 的教程和示例非常細緻，並且也都以 ipynb 格式提供，大幅降低 GoMLX 上手門檻。

不過，儘管 GoMLX 展示了強大的潛力，但目前它仍處於初始階段，許多核心功能尚未完全實現。但總體而言，GoMLX 代表了 Go 語言在機器學習領域的一次重要嘗試，其設計理念和目標值得關注。儘管目前還存在諸多挑戰和未完成的工作，但我們可以對其未來的發展保持謹慎樂觀的態度。

如果你有機器學習開發經驗，不妨參與到 GoMLX 的開發中，爲這一項目貢獻力量，共同推動 Go 在機器學習領域的進步。

參考資料

[1] 

Google 前研究員 Jan Pfeifer: https://github.com/janpfeifer

[2] 

Go 機器學習框架項目 GoMLX: https://github.com/gomlx/gomlx

[3] 

GoMLX 教程: https://gomlx.github.io/gomlx/notebooks/tutorial.html

[4] 

GoNB 項目: https://github.com/janpfeifer/gonb

[5] 

Gopher 部落知識星球: https://public.zsxq.com/groups/51284458844544

[6] 

鏈接地址: https://m.do.co/c/bff6eed92687

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/kFlsQ7v84dLrnNPz8Ktczw