PyTorch還是TensorFlow？這有一份新手深度學習框架選擇指南

01-26

問耕編譯整理

量子位出品 | 公眾號 QbitAI

前幾天，量子位發過一篇《忽悠VC指南》。其中有一條建議是，當你假裝AI專家時，最好別談眾人皆知的TensorFlow，那談什麼？

PyTorch。

其實這也不全然都是調侃。和TensorFlow相比，PyTorch確實也有為數不少的擁躉。根據Keras作者Fran?ois Chollet發布的數據顯示：GitHub上的各種深度學習框架中，PyTorch排在第五位。

第二季度深度學習框架排名：

但需要提醒大家的是，PyTorch可是今年1月19日才正式發布。

最近，斯坦福大學研究機器學習的博士生Awni Hannun，圍繞PyTorch還是TensorFlow這個話題，做了一個深入的比較。量子位把內容傳送如下：

我寫的這份指南，主要對比了PyTorch和TensorFlow之間的區別。希望對那些想著手開始一個新項目或者考慮轉換深度學習框架的人有所幫助。

我主要考察了深度學習堆棧的訓練和部署時，相關的可編程性和靈活性。在這裡，我不會過多的談論速度、內存佔用等方面的權衡。

先放結論

PyTorch更有利於研究人員、愛好者、小規模項目等快速搞出原型。而TensorFlow更適合大規模部署，特別是需要跨平台和嵌入式部署時。

然後咱們一項一項分著說。

上手時間

贏家：PyTorch

PyTorch本質上是Numpy的替代者，而且支持GPU、帶有高級功能，可以用來搭建和訓練深度神經網路。如果你熟悉Numpy、Python以及常見的深度學習概念（卷積層、循環層、SGD等），會非常容易上手PyTorch。

而TensorFlow可以看成是一個嵌入Python的編程語言。你寫的TensorFlow代碼會被Python編譯成一張圖，然後由TensorFlow執行引擎運行。我見過好多新手，因為這個增加的間接層而困擾。也正是因為同樣的原因，TensorFlow有一些額外的概念需要學習，例如會話、圖、變數作用域（variable scoping）、佔位符等。

另外還需要更多的樣板代碼才能讓一個基本的模型運行。所以TensorFlow的上手時間，肯定要比PyTorch長。

圖創建和調試

贏家：PyTorch

創建和運行計算圖可能是兩個框架最不同的地方。在PyTorch中，圖結構是動態的，這意味著圖在運行時構建。而在TensorFlow中，圖結構是靜態的，這意味著圖先被「編譯」然後再運行。

舉一個簡單的例子，在PyTorch中你可以用標準的Python語法編寫一個for循環結構

for _ in range(T): h = torch.matmul(W, h) + b

此處T可以在每次執行代碼時改變。而TensorFlow中，這需要使用「控制流操作」來構建圖，例如tf.while_loop。TensorFlow確實提供了dynamic_rnn用於常見結構，但是創建自定義動態計算真的更加困難。

PyTorch中簡單的圖結構更容易理解，更重要的是，還更容易調試。調試PyTorch代碼就像調試Python代碼一樣。你可以使用pdb並在任何地方設置斷點。調試TensorFlow代碼可不容易。要麼得從會話請求要檢查的變數，要麼學會使用TensorFlow的調試器（tfdbg）。

全面性

贏家：TensorFlow

隨著PyTorch逐漸成熟，我預計這部分的差距會趨近於零。但目前，TensorFlow還是有一些PyTorch不支持的功能。它們是：

沿維翻轉張量（np.flip, np.flipud, np.fliplr）
檢查無窮與非數值張量（np.is_nan, np.is_inf）
快速傅里葉變換（np.fft）

這些TensorFlow都支持。另外，TensorFlow的contrib軟體包中，有更多PyTorch沒有的高級功能和模型。

序列化

贏家：TensorFlow

兩種框架下保存和載入模型都很簡單。PyTorch有一個特別簡單的API，可以保存模型的所有權重或pickle整個類。TensorFlow的Saver對象也很易用，而且為檢查提供了更多的選項。

TensorFlow序列化的主要優點是可以將整個圖保存為protocol buffer。包括參數和操作。然而圖還能被載入進其他支持的語言（C++、Java）。這對於部署堆棧至關重要。理論上，當你想改動模型源代碼但仍希望運行舊模型時非常有用。

部署

贏家：TensorFlow

對於小規模的伺服器端部署（例如一個Flask web server），兩個框架都很簡單。

對於移動端和嵌入式部署，TensorFlow更好。不只是比PyTorch好，比大多數深度學習框架都要要。使用TensorFlow，部署在Android或iOS平台時只需要很小的工作量，至少不必用Java或者C++重寫模型的推斷部分。

對於高性能伺服器端的部署，還有TensorFlow Serving能用。除了性能之外，TensorFlow Serving一個顯著的優點是可以輕鬆的熱插拔模型，而不會使服務失效。

文檔

贏家：平手

對於兩個框架，我都在文檔中找到所需的一切。Python API被很好的記錄，以及有足夠的案例和教程來學習框架。

一個特例是，PyTorch的C庫大多數沒有文檔。不過，這隻有在你編寫一個定製化的C擴展時才有影響。

數據載入

贏家：PyTorch

PyTorch中用於載入數據的API設計的很棒。介面由一個數據集、一個取樣器和一個數據載入器構成。數據載入器根據取樣器的計劃，基於數據集產生一個迭代器。並行化數據載入簡單的就像把num_workers參數傳遞給數據載入器一樣簡單。

我在TensorFlow中沒有發現特別有用的數據載入工具。很多時候，並不總能直接把準備並行運行的預處理代碼加入TensorFlow圖。以及API本身冗長難學。

設備管理

贏家：TensorFlow

TensorFlow的設備管理非常好用。通常你不需要進行調整，因為默認的設置就很好。例如，TensorFlow會假設你想運行在GPU上（如果有的話）。而在PyTorch中，即使啟用了CUDA，你也需要明確把一切移入設備。

TensorFlow設備管理唯一的缺點是，默認情況下，它會佔用所有的GPU顯存。簡單的解決辦法是指定CUDA_VISIBLE_DEVICES。有時候大家會忘了這一點，所以GPU在空閑的時候，也會顯得很忙。

在PyTorch中，我發現代碼需要更頻繁的檢查CUDA是否可用，以及更明確的設備管理。在編寫能夠同時在CPU和GPU上運行的代碼時尤其如此。以及得把GPU上的PyTorch變數轉換為Numpy數組，這就顯得有點冗長。

numpy_var = variable.cpu().data.numpy()

自定義擴展

贏家：PyTorch

兩個框架都可以構建和綁定用C、C++、CUDA編寫的自定義擴展。TensorFlow仍然需要更多的樣板代碼，儘管這對於支持多類型和設備可能更好。在PyTorch中，你只需為每個CPU和GPU編寫一個介面和相應的實現。兩個框架中編譯擴展也是直接記性，並不需要在pip安裝的內容之外下載任何頭文件或者源代碼。

關於TensorBoard

TensorBoard是TensorFlow自帶的可視化工具，用來查看機器學習訓練過程中數據的變化。通過訓練腳本中的幾個代碼段，你可以查看任何模型的訓練曲線和驗證結果。TensorBoard作為web服務運行，特別便於對於無頭結點上存儲的結果進行可視化。

如果沒有類似的功能，我可不想用PyTorch。不過還好，藉助兩個開源項目可以實現。第一個是tensorboard_logger，第二個是crayon。tensorboard_logger庫用起來甚至比TensorBoard的「摘要」更容易，儘管想用這個首先得安裝TensorBoard。

crayon可以完全替代TensorBoard，但是需要更多的設置（docker是先決條件）。

關於Keras

Keras是具有可配置後端的高級API。目前TensorFlow、Theano、CNTK都支持。也許不久的將來，PyTorch也會提供支持。作為tf.contrib的一部分，Keras也隨TensorFlow一起分發。

雖然上面我沒有討論過Keras，但其API特別容易使用，這也是配合常見深度神經網路架構最快的方式。不過，使用API畢竟沒有使用PyTorch或者核心TensorFlow那麼靈活。

Keras是許多常用的深層神經網路架構中運行最快的方法之一。

關於TensorFlow Fold

今年2月，谷歌推出了TensorFlow Fold。這個庫建立在TensorFlow智商，允許構建更多的動態圖。這個庫的主要優勢是動態批處理。動態批處理可以自動對不同大小的輸入進行批量計算（例如解析樹上的循環網路）。

可編程性方面，語法不像PyTorch那麼簡單，當然在某些情況下，批處理帶來的性能提升還是值得考慮。

好了，以上就是來自Awni Hannun的分享。希望對大家有所幫助，另外兼聽則明，量子位繼續搬運幾條reddit上讀者的評論。

讀者Reiinakano：

我不覺得「文檔」環節雙方平分秋色。
官方的TensorFlow文檔很差勁。比方，在PyTorch文檔中，對於遷移學習的講解，使用了實際、有用的代碼，而且還解釋了構建的方式。而在TensorFlow的文旦中，整個講解就是運行了一個bash scripts，沒有任何實際代碼。

讀者ThaHypnotoad：

PyTorch還有很長的路要走。前幾天我發現int tensor沒有neg()定義。不過，我更討厭調試靜態圖。所以我選擇忍受PyTorch成長的煩惱。

讀者trias10：

另一個缺點是，PyTorch沒有對Windows的官方支持，而TensorFlow有。很多場合（通常是金融和保險行業）使用Windows進行開發和原型設計，用Linux進行生產部署，所以你需要一個能對兩者都支持的框架。

好啦，今天先談到這裡。希望對你有所幫助~

— 完 —

歡迎大家關注我們的專欄：量子位 - 知乎專欄

誠摯招聘

量子位正在招募編輯/記者，工作地點在北京中關村。期待有才氣、有熱情的同學加入我們！相關細節，請在量子位公眾號(QbitAI)對話界面，回復「招聘」兩個字。

量子位 QbitAI

?"?" ? 追蹤AI技術和產品新動態