放棄TensorFlow，谷歌全面轉向JAX

OpenCV學堂 2022-06-15 16:00

NFC 還能這樣用？瑞薩無線充電方案全解析 一站式破解芯片研發全流程痛點，西門子 EDA Forum 2026等你來

點擊上方↑↑↑“OpenCV學堂”關注我

來源：公眾號?機器之心??授權

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 團隊。

論文地址：https://cs.stanford.edu/~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 有一天也會被新框架取代，這尚未可知。但至少，新框架出現的時機已經成熟。

原文鏈接：https://www.businessinsider.com/facebook-pytorch-beat-google-tensorflow-jax-meta-ai-2022-6