放棄TensorFlow，谷歌全面轉向JAX

TensorFlow 大概已經(jīng)成為了谷歌的一枚「棄子」。

2015 年，谷歌大腦開放了一個名為「TensorFlow」的研究項目，這款產品迅速流行起來，成為人工智能業(yè)界的主流深度學習框架，塑造了現(xiàn)代機器學習的生態(tài)系統(tǒng)。

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

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

近日，外媒 Business Insider 采訪了一系列開發(fā)人員、硬件專家、云供應商以及與谷歌機器學習工作關系密切的人，獲得了同樣的觀點。TensorFlow 已經(jīng)輸?shù)袅诉@場戰(zhàn)爭，其中有人用了一個鮮明的比喻：「PyTorch 吃掉了 TensorFlow 的午餐?！?br />
專家們表示，鑒于戰(zhàn)術失誤、開發(fā)決策和 Meta 在開源社區(qū)中的一系列智取策略，谷歌引領互聯(lián)網(wǎng)機器學習未來的機會正在逐漸消失。

在 PyTorch 的陰影下，谷歌正在悄悄地開發(fā)一個機器學習框架，就是 JAX（曾是「Just After eXecution」的首字母縮寫，但官方說法中不再代表任何東西），許多人將其視為 TensorFlow 的繼承者。

接近該項目的人士告訴 Insider，谷歌大腦和 DeepMind 在很大程度上放棄了 TensorFlow，轉而使用 JAX。這為谷歌的其他部門鋪平了跟隨的道路，一位谷歌內部人士向 Insider 證實，JAX 現(xiàn)在幾乎已在谷歌大腦和 DeepMind 中被全球采用。

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

熟悉該項目的人士表示，Jax 現(xiàn)在有望成為未來幾年所有使用機器學習的谷歌產品的支柱，就像 TensorFlow 在 2015 年之后幾年所做的那樣。

「JAX 是一項工程壯舉，」Julia 編程語言創(chuàng)建者 Viral Shah 說?！肝艺J為 JAX 是一種通過 Python 實例化的獨立編程語言。如果你遵守 JAX 想要的規(guī)則，它就可以發(fā)揮它的魔力，這真是令人驚嘆?！?br />
現(xiàn)在，谷歌希望在這場競賽中再次獲得金牌，同時也從開發(fā) TensorFlow 時所犯的錯誤中吸取教訓，但這將是一個巨大的挑戰(zhàn)。

TensorFlow 的暮光，PyTorch 的崛起

根據(jù)提供給 Insider 的數(shù)據(jù)，PyTorch 在一些必讀開發(fā)者論壇上的帖子正在迅速趕超 TensorFlow。Stack Overflow 的參與度數(shù)據(jù)顯示，以論壇問題份額衡量， TensorFlow 受歡迎程度近年來停滯不前，而 PyTorch 的參與度繼續(xù)上升。


TensorFlow 起步強勁，推出后受歡迎程度不斷提高。Uber 和 Airbnb 等公司以及 NASA 等機構很快就開始將其用于一些復雜的項目，這些項目需要在大量數(shù)據(jù)集上訓練算法。截至 2020 年 11 月，TensorFlow 已被下載 1.6 億次。

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

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

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

「我們主要使用 PyTorch，它擁有最多的社區(qū)支持，」機器學習初創(chuàng)公司 Hugging Face 的研究工程師 Patrick von Platten 說?！肝覀冋J為 PyTorch 可能在開源方面做得最好，他們能確保在線回復問題，所有示例都能 work?！?br />
一些最大的組織機構開始在 PyTorch 上運行項目，包括那些曾經(jīng)依賴 TensorFlow 的機構。不久之前，特斯拉、Uber 等公司就在 PyTorch 上運行了他們最艱巨的機器學習研究項目。

TensorFlow 的新增功能有時會復制使 PyTorch 流行的元素，使得 TensorFlow 對于其最初的研究人員和用戶受眾來說越來越臃腫。一個這樣的例子是它在 2017 年增加了「敏銳執(zhí)行」，這是 Python 的原生特性，使開發(fā)人員可以輕松分析和調試他們的代碼。

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

團隊研究人員 Roy Frostige、Matthew James Johnson 和 Chris Leary 在 2018 年發(fā)表了一篇名為《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 在很多方面仍然依賴于其他框架。開發(fā)人員表示：JAX 沒有加載數(shù)據(jù)和預處理數(shù)據(jù)的簡單方法，需要使用 TensorFlow 或 PyTorch 進行大量數(shù)據(jù)處理。

JAX 基礎框架 XLA 也針對谷歌 TPU 設備做了很大的優(yōu)化。該框架還適用于更傳統(tǒng)的 GPU 和 CPU。了解該項目的人稱：該項目仍有辦法對 GPU 和 CPU 進行優(yōu)化，以達到 TPU 同等水平。

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

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

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

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