譯者 | 朱先忠

審校 | 墨色

深度學習每天都在以各種形式影響著我們的生活。無論是基于用戶語音命令的Siri、Alexa、手機上的實時翻譯應(yīng)用程序，還是支持智能拖拉機、倉庫機器人和自動駕駛汽車的計算機視覺技術(shù)，每個月似乎都會迎來新的進展。幾乎所有這些深度學習應(yīng)用程序的編寫都來自于這三種框架：TensorFlow、PyTorch或者JAX。

那么，你到底應(yīng)該使用哪些深度學習框架呢？在本文中，我們將對TensorFlow、PyTorch和JAX進行高級比較。我們的目標是讓你了解發(fā)揮其優(yōu)勢的應(yīng)用程序類型，當然還要考慮社區(qū)支持和易用性等因素。

你應(yīng)該使用TensorFlow嗎？

“從來沒有人因為購買IBM而被解雇”是20世紀70年代和80年代計算機界的口號。在本世紀初，使用TensorFlow進行深度學習也是如此。但眾所周知，進入20世紀90年代時，IBM就已被“擱置一旁”。那么，TensorFlow在2015年首次發(fā)布后7年的今天以及未來新的十年中仍然具有競爭力嗎？

當然。TensorFlow并不是一直都在原地踏步。首先，TensorFlow 1.x是以一種非Python的方式構(gòu)建靜態(tài)圖的，但是在TensorFlow 2.x中，還可以使用動態(tài)圖模式（eager mode）構(gòu)建模型，以便立即評估操作，這讓人感覺它更像PyTorch。在高層，TensorFlow提供了Keras以便于開發(fā)；在底層，它提供了XLA（Accelerated Linear Algebra，加速線性代數(shù)）優(yōu)化編譯器以提高速度。XLA在提高GPU性能方面發(fā)揮了神奇作用，它是利用谷歌TPU（Tensor Processing Units，張量處理單元）能力的主要方法，為大規(guī)模模型訓練提供了無與倫比的性能。

其次，多年來TensorFlow一直努力盡可能在所有方面做得很好。例如，你是否想要在成熟的平臺上以定義良好且可重復(fù)的方式為模型提供服務(wù)？TensorFlow隨時可以提供服務(wù)。你是否想要將模型部署重新定位到web、智能手機等低功耗計算或物聯(lián)網(wǎng)等資源受限設(shè)備？在這一點上，TensorFlow.js和TensorFlow Lite都已經(jīng)非常成熟。

顯然，考慮到Google仍然在百分之百地使用TensorFlow運行其生產(chǎn)部署，就可以確信TensorFlow一定能夠應(yīng)用戶的規(guī)模需求。

但是，近來確實有一些項目中的因素不容忽視。簡而言之，把項目從TensorFlow 1.x升級到TensorFlow 2.x其實是非常殘酷的。一些公司考慮到更新代碼后在新的版本上正常工作所需的努力，干脆決定將代碼移植到PyTorch框架下。此外，TensorFlow在科研領(lǐng)域也失去了動力，幾年前已開始傾向于PyTorch提供的靈活性，這導致TensorFlow在研究論文中的使用不斷減少。

此外，“Keras事件”也沒有起到任何作用。Keras在兩年前成為TensorFlow發(fā)行版的一個集成部分，但最近又被拉回到一個單獨的庫中，并確定了自己的發(fā)行計劃。當然，排除Keras不會影響開發(fā)人員的日常生活，但在框架的一個小更新版本中出現(xiàn)如此引人注目的變化，并不會激發(fā)程序員使用TensorFlow框架的信心。

話雖如此，TensorFlow的確還是一個可靠的框架，它擁有廣泛的深度學習生態(tài)系統(tǒng)，使用者可以在TensorFlow上構(gòu)建適用于所有規(guī)模的應(yīng)用程序和模型。如果真的這樣做，將會有很多不錯的合作公司。但如今，TensorFlow可能還不是首選。

你應(yīng)該使用PyTorch嗎？

PyTorch不再是緊跟TensorFlow之后的“新貴”，而是當今深度學習領(lǐng)域的主要力量，可能主要用于研究，但也越來越多地用于生產(chǎn)應(yīng)用。隨著動態(tài)圖模式（eager mode）成為TensorFlow和PyTorch中開發(fā)的默認方法，PyTorch的自動微分（autograd）提供的更具Python風格的方法似乎贏得了與靜態(tài)圖的戰(zhàn)爭。

與TensorFlow不同的是，自0.4版本中不推薦使用變量API以來，PyTorch的核心代碼沒有經(jīng)歷過任何重大的中斷。以前，變量需要使用自動生成張量，而現(xiàn)在，一切都是張量。但這并不是說無論在哪兒都不存在錯誤。例如，如果你一直在使用PyTorch跨多個GPU進行訓練，可能會遇到DataParallel和較新的DistributedDataParaller之間的差異。你應(yīng)該經(jīng)常使用DistributedDataParallel，但實際上并不反對使用DataParaller。

雖然PyTorch在XLA/TPU支持方面一直落后于TensorFlow和JAX，但截至2022年，情況已經(jīng)有了很大改善。PyTorch現(xiàn)在支持訪問TPU虛擬機，支持老式TPU節(jié)點支持，以及支持在CPU、GPU或TPU上運行代碼的簡單命令行部署，而無需更改代碼。如果你不想處理PyTorch經(jīng)常讓你編寫的一些樣板代碼，那么你可以求助于Pytorche Lightning這樣更高級別的擴展程序，它讓你專注于實際工作，而不是重寫訓練循環(huán)。而另一方面，雖然PyTorch Mobile的工作仍在繼續(xù)，但它遠不如TensorFlow Lite那么成熟。

在生產(chǎn)方面，PyTorch現(xiàn)在可以與Kubeflow等框架無關(guān)平臺進行集成，而且TorchServe項目可以處理擴展、度量和批量推理等部署細節(jié)——在PyTorch開發(fā)人員自己維護的小軟件包中能夠提供所有MLOps優(yōu)點。另一方面，PyTorch支持規(guī)?？s放嗎？沒有問題！Meta公司多年來一直在生產(chǎn)領(lǐng)域運行PyTorch；所以，任何人告訴你PyTorch無法處理大規(guī)模的工作負載其實都是謊言。盡管如此，有一種情況是，PyTorch可能不像JAX那樣友好，特別是在需要大量GPU或TPU進行非常大量的訓練方面。

最后，依然存在一個人們不愿提及的棘手問題——PyTorch在過去幾年的受歡迎程度幾乎離不開Hugging Face公司的Transformers庫的成功。是的，Transformers現(xiàn)在也支持TensorFlow和JAX，但它最初是一個PyTorch項目，仍然與框架緊密結(jié)合。隨著Transformer架構(gòu)的興起，PyTorch在研究方面的靈活性，以及通過Hugging Face的模型中心在發(fā)布后幾天或幾個小時內(nèi)引入如此多的新模型的能力，很容易看出為什么PyTorch在這些領(lǐng)域如此流行。

你應(yīng)該使用JAX嗎？

如果你對TensorFlow不感興趣，那么Google可能會為你提供其他服務(wù)。JAX是一個由Google構(gòu)建、維護和使用的深度學習框架，但它不是官方的Google產(chǎn)品。然而，如果你留意過去一年左右Google/DeepMind的論文和產(chǎn)品發(fā)布，你就會注意到Google的許多研究已經(jīng)轉(zhuǎn)移到了JAX。因此，盡管JAX并不是谷歌的“官方”產(chǎn)品，但它是谷歌研究人員用來推動邊界的東西。

到底什么是JAX呢？理解JAX的一個簡單方法是：想象一個GPU/TPU加速版本的NumPy，它可以用“一根魔杖”神奇地將Python函數(shù)矢量化，并處理所有這些函數(shù)的導數(shù)計算。最后，它提供了一個即時（JIT：Just-In-Time）組件，用于獲取代碼并為XLA（Accelerated Linear Algebra，即加速線性代數(shù)）編譯器進行優(yōu)化，從而大幅提高TensorFlow和PyTorch的性能。目前一些代碼的執(zhí)行速度提高了四到五倍，只需在JAX中重新實現(xiàn)，而不需要進行任何真正的優(yōu)化工作。

考慮到JAX是在NumPy級別工作的，JAX代碼是在比TensorFlow/Keras（甚至是PyTorch）低得多的級別上編寫的。令人高興的是，有一個小型但不斷增長的生態(tài)系統(tǒng)，圍繞著JAX進行了一些擴展。你想要使用神經(jīng)網(wǎng)絡(luò)庫嗎？當然可以。其中有來自谷歌的Flax，還有來自DeepMind（也包括谷歌）的Haiku。此外，Optax可滿足你的所有優(yōu)化器需求，PIX可用于圖像處理，此外還有更多功能。一旦你使用Flax之類的東西，構(gòu)建神經(jīng)網(wǎng)絡(luò)就變得相對容易掌握。請注意，仍然有一些略讓人糾結(jié)的問題。例如，經(jīng)驗豐富的人經(jīng)常談到JAX處理隨機數(shù)的方式與許多其他框架不同。

那么，你是否應(yīng)該將所有內(nèi)容轉(zhuǎn)換為JAX并利用這一前沿技術(shù)呢？這個問題因人而異。如果你深入研究需要大量資源來訓練的大規(guī)模模型的話，建議采用這種方法。此外，如果你關(guān)注JAX在確定性訓練，以及其他需要數(shù)千個TPU Pod的項目，那么，也值得一試。?

小結(jié)

因此，結(jié)論是什么呢？你應(yīng)該使用哪種深度學習框架？遺憾的是，這道題并沒有唯一的答案，完全取決于你正在處理的問題類型、計劃部署模型以處理的規(guī)模，甚至還依賴于你所面對的計算平臺。

不過，如果你從事的是文本和圖像領(lǐng)域，并且正在進行中小型研究，以期在生產(chǎn)中部署這些模型，那么PyTorch可能是目前最好的選擇。從最近的版本看，它正好針對這類應(yīng)用空間的最佳點。

如果你需要從低計算設(shè)備中獲取所有性能，那么建議你使用TensorFlow以及極為堅固的TensorFlow Lite軟件包。最后，如果你正在研究數(shù)百億、數(shù)千億或更多參數(shù)的訓練模型，并且你主要是為了研究目的而訓練它們，那么也許是時候試一試JAX了。

原文鏈接：https://www.infoworld.com/article/3670114/tensorflow-pytorch-and-jax-choosing-a-deep-learning-framework.html

譯者介紹

朱先忠，51CTO社區(qū)編輯，51CTO專家博客、講師，濰坊一所高校計算機教師，自由編程界老兵一枚。