放棄TensorFlow，谷歌全面轉向JAX

TensorFlow 大概已經成為了谷歌的一枚「棄子」。

2015 年，谷歌大腦開放了一個名為「TensorFlow」的研究專案，這款產品迅速流行起來，成為人工智慧業界的主流深度學習框架，塑造了現代機器學習的生態系統。

七年後的今天，故事的走向已經完全不同：谷歌的 TensorFlow 失去了開發者的擁護，因為他們已經轉向了 Meta 推出的另一款框架 PyTorch。

曾經無處不在的機器學習工具 TensorFlow 已經悄悄落後，而 Facebook 在 2017 年開源的 PyTorch 正在成為這個領域的霸主。

近日，外媒 Business Insider 採訪了一系列開發人員、硬體專家、雲供應商以及與谷歌機器學習工作關係密切的人，獲得了同樣的觀點。TensorFlow 已經輸掉了這場戰爭，其中有人用了一個鮮明的比喻：「PyTorch 吃掉了 TensorFlow 的午餐。」

專家們表示，鑑於戰術失誤、開發決策和 Meta 在開源社群中的一系列智取策略，谷歌引領網際網路機器學習未來的機會正在逐漸消失。

在 PyTorch 的陰影下，谷歌正在悄悄地開發一個機器學習框架，就是 JAX（曾是「Just After eXecution」的首字母縮寫，但官方說法中不再代表任何東西），許多人將其視為 TensorFlow 的繼承者。 接近該專案的人士告訴 Insider，谷歌大腦和 DeepMind 在很大程度上放棄了 TensorFlow，轉而使用 JAX。這為谷歌的其他部門鋪平了跟隨的道路，一位谷歌內部人士向 Insider 證實，JAX 現在幾乎已在谷歌大腦和 DeepMind 中被全球採用。

接近谷歌機器學習工作的人士表示，最初 JAX 面臨著來自內部的強烈反對，一些人認為谷歌員工已經習慣了使用 TensorFlow。儘管它可能很難用，但它一直是谷歌員工中的統一因素。他們說，JAX 方法是要簡單得多，但它會改變 Google 內部構建軟體的方式。

熟悉該專案的人士表示，Jax 現在有望成為未來幾年所有使用機器學習的谷歌產品的支柱，就像 TensorFlow 在 2015 年之後幾年所做的那樣。

「JAX 是一項工程壯舉，」Julia 程式語言建立者 Viral Shah 說。「我認為 JAX 是一種透過 Python 例項化的獨立程式語言。如果你遵守 JAX 想要的規則，它就可以發揮它的魔力，這真是令人驚歎。」

現在，谷歌希望在這場競賽中再次獲得金牌，同時也從開發 TensorFlow 時所犯的錯誤中吸取教訓，但這將是一個巨大的挑戰。

TensorFlow 的暮光，PyTorch 的崛起

根據提供給 Insider 的資料，PyTorch 在一些必讀開發者論壇上的帖子正在迅速趕超 TensorFlow。Stack Overflow 的參與度資料顯示，以論壇問題份額衡量， TensorFlow 受歡迎程度近年來停滯不前，而 PyTorch 的參與度繼續上升。 TensorFlow 起步強勁，推出後受歡迎程度不斷提高。Uber 和 Airbnb 等公司以及 NASA 等機構很快就開始將其用於一些複雜的專案，這些專案需要在大量資料集上訓練演算法。截至 2020 年 11 月，TensorFlow 已被下載 1.6 億次。

但谷歌持續及增量的功能更新使得 TensorFlow 變得笨拙，且對使用者不友好，即使是谷歌內部的那些人、開發人員和與專案關係密切的人都認為如此。隨著機器學習領域以驚人的速度發展，谷歌不得不經常使用新工具更新其框架。接近該專案的人士表示，該專案已經在內部傳播開來，越來越多的人參與其中，不再專注最初是什麼讓 TensorFlow 成為首選工具。

專家告訴 Insider，對於許多擁有引領者身份的公司來說，這種瘋狂的貓鼠遊戲是一個反覆出現的問題。例如，谷歌並不是第一家建立搜尋引擎的公司，它能夠從 AltaVista 或 Yahoo 等前輩的錯誤中吸取教訓。

2018 年，PyTorch 推出了完整版。雖然 TensorFlow 和 PyTorch 都建立在 Python 之上，但 Meta 在滿足開源社群的需求方面投入了大量資金。熟悉 TensorFlow 專案的人士說，PyTorch 還受益於專注做一些 TensorFlow 團隊錯過的事情。

「我們主要使用 PyTorch，它擁有最多的社群支援，」機器學習初創公司 Hugging Face 的研究工程師 Patrick von Platten 說。「我們認為 PyTorch 可能在開源方面做得最好，他們能確保線上回覆問題，所有示例都能 work。」

一些最大的組織機構開始在 PyTorch 上執行專案，包括那些曾經依賴 TensorFlow 的機構。不久之前，特斯拉、Uber 等公司就在 PyTorch 上執行了他們最艱鉅的機器學習研究專案。

TensorFlow 的新增功能有時會複製使 PyTorch 流行的元素，使得 TensorFlow 對於其最初的研究人員和使用者受眾來說越來越臃腫。一個這樣的例子是它在 2017 年增加了「敏銳執行」，這是 Python 的原生特性，使開發人員可以輕鬆分析和除錯他們的程式碼。

嘗試用 JAX 自救

隨著 PyTorch 和 TensorFlow 之間競爭日益激烈，谷歌內部的一個小型研究團隊開發了一個新框架 JAX，該框架將更容易訪問張量處理單元（TPU）——一種谷歌專門為機器學習和 TensorFlow 定製的晶片。

團隊研究人員 Roy Frostige、Matthew James Johnson 和 Chris Leary 在 2018 年發表了一篇名為《Compilation of machine learning software through high-level traceability》的論文，介紹了這個新框架 JAX。PyTorch 的原始作者之一 Adam Paszky 於 2020 年初全職加入 JAX 團隊。 論文地址：~rfrostig/pubs/jax-mlsys2018.pdf

JAX 提出了一種直接的設計來解決機器學習中最複雜的問題之一：將一個大問題的工作分散到多個晶片上。JAX 不會為不同的晶片執行單段程式碼，而是自動分配工作，即時訪問儘可能多的 TPU，以滿足執行需要。這解決了谷歌研究人員對算力的巨大需求。

PyTorch 和 TensorFlow 都是以同樣的方式開始的，首先是研究專案，然後成為機器學習研究的標準化工具，從學界擴散到更多地方。

然而，JAX 在很多方面仍然依賴於其他框架。開發人員表示：JAX 沒有載入資料和預處理資料的簡單方法，需要使用 TensorFlow 或 PyTorch 進行大量資料處理。

JAX 基礎框架 XLA 也針對谷歌 TPU 裝置做了很大的最佳化。該框架還適用於更傳統的 GPU 和 CPU。瞭解該專案的人稱：該專案仍有辦法對 GPU 和 CPU 進行最佳化，以達到 TPU 同等水平。

谷歌發言人表示，2018 年至 2021 年穀歌內部對 GPU 的支援策略欠優，缺乏與大型 GPU 供應商英偉達的合作，因此谷歌轉而關注 TPU，谷歌自己內部的研究也主要集中在 TPU 上，導致缺乏良好的 GPU 使用反饋迴圈。從這個意義上說，谷歌對 TPU 的重視和專注也屬迫不得已。

然而，如 Cerebras Systems 的 CEO Andrew Feldman 所說：「任何以一種裝置優於另一種裝置的行為都是不良做法，會遭到開源社群的抵制。沒有人希望侷限於單一的硬體供應商，這就是機器學習框架應運而生的原因。機器學習從業者希望確保他們的模型是可移植的，可以移植到他們選擇的任何硬體上，而不是被鎖定在一個平臺上。」

如今，PyTorch 快 6 歲了，TensorFlow 在這個年紀早已出現衰落的跡象。也許 PyTorch 有一天也會被新框架取代，這尚未可知。但至少，新框架出現的時機已經成熟。