PyTorch和TensorFlow：九項對比讀各自長項短板

選自GitHub

作者：Awni Hannun

機器之心編譯

參與：Panda

現在是各種機器學習框架群雄爭霸的時代，各種各樣的比較文章也層出不窮。近日，斯坦福大學計算機科學系博士生 Awni Hannun 也發表了一篇文章，談了自己對 PyTorch 和 TensorFlow 這兩大明星框架的心得體驗，並在不同的方面對這兩者進行了比較，機器之心對本文進行了編譯介紹。

這篇指南主要介紹了我找到的 PyTorch 和 TensorFlow 之間的不同之處。這篇文章的目的是幫助那些想要開始一個新項目或從一種深度學習框架切換到另一種框架的人。本文重點關注的是在設置訓練組件和部署深度學習時的可編程性和靈活性。我不會深入到性能方面（速度/內存佔用）的比較。

概要

PyTorch 更適用於研究、愛好者和小規模項目的快速原型開發。TensorFlow 更適合大規模部署，尤其是涉及跨平台和嵌入式部署時。

上手時間

獲勝者：PyTorch

PyTorch 本質上是支持 GPU 的 NumPy 替代，配備了可用於構建和訓練深度神經網路的更高級的功能。所以如果你熟悉 NumPy、Python 和常用的深度學習抽象（卷積層、循環層、SGD 等），那 PyTorch 就很容易學。

另一方面，則可以將 TensorFlow 看作是一種嵌入 Python 的編程語言。當你編寫 TensorFlow 代碼時，它會被 Python「編譯」成圖（graph），然後由 TensorFlow 執行引擎運行。我看到過有些 TensorFlow 新手難以理解這額外增加的間接一層工序。同樣因為這個原因，TensorFlow 還有一些需要額外學習的概念，比如會話（session）、圖、變數範圍和佔位符。要讓基本的模型跑起來也需要更多樣板代碼。上手 TensorFlow 的時間肯定會比 PyTorch 長。

創建圖和調試

獲勝者：PyTorch

創建和運行計算圖可能是這兩個框架差別最大的地方。在 PyTorch 中，圖結構是動態的，也就是說圖是在運行時創建的。在 TensorFlow 中，圖結構是靜態的，也就是說圖在「編譯」之後再運行。舉個簡單例子，在 PyTorch 中，你可以使用標準的 Python 句法寫一個 for 循環：

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

而且 T 可以在這段代碼的執行之間改變。在 TensorFlow 中，這需要在構建圖時使用控制流操作（control flow operations），比如 tf.while_loop。TensorFlow 確實有 dynamic_rnn 可用於更常見的結構，但創建自定義的動態計算也更加困難。

PyTorch 簡單的圖構建方式更容易理解，但也許更重要的是也更容易調試。調試 PyTorch 代碼就跟調試 Python 代碼一樣。你可以使用 pdb，並且可以在任何地方設置斷點。調試 TensorFlow 則沒這麼容易。它有兩個選擇，一是從會話中請求你想檢查的變數，而是學會使用 TensorFlow 調試器（tfdbg）。

覆蓋度

獲勝者：TensorFlow

隨著 PyTorch 的發展，我預計這兩者之間的差距會縮小至零。但是，TensorFlow 仍然支持一些 PyTorch 並不支持的功能。PyTorch 目前還不具備的特性包括：

• 沿維度方向的張量翻轉（np.flip、 np.flipud、 np.fliplr）

• 檢查張量是否為 NaN 和無窮大（np.is_nan、np.is_inf）

• 快速傅立葉變換（np.fft）

而 TensorFlow 支持所有這些。另外比起 PyTorch，TensorFlow 的 contrib 包也有遠遠更多更高級的函數和模型。

序列化（serialization）

獲勝者：TensorFlow

在這兩種框架中，保存和載入模型都很簡單。PyTorch 有一個非常簡單的 API，既可以保存模型的所有權重，也可以 pickle（加工）整個類。TensorFlow 的 Saver 對象也很容易使用，而且也為檢查點提供了更多選擇。

TensorFlow 在序列化方面的主要優勢是整個計算圖都可以保存為 protocol buffer。這既包括參數，也包括運算。然後這個圖可以用其它支持的語言（C++、Java）載入。對於不支持 Python 的部署環境來說，這是非常重要的功能。而且理論上，這個功能也可以在你修改模型的源代碼，但又想運行舊模型時為你提供幫助。

部署

獲勝者：TensorFlow

對於小型伺服器（比如 Flask 網頁伺服器）上的部署，兩種框架都很簡單。

TensorFlow 支持移動和嵌入式部署，而包括 PyTorch 在內的很多深度學習框架都沒有這個能力。在 TensorFlow 上，要將模型部署到安卓或 iOS 上需要不小的工作量，但至少你不必使用 Java 或 C++ 重寫你模型的整個推理部分。

對於高性能伺服器上的部署，還有 TensorFlow Serving 可用。我還沒體驗過 TensorFlow Serving，所以我不能說它有哪些優缺點。對於嚴重依賴機器學習的服務，我猜想 TensorFlow Serving 可能就是繼續使用 TensorFlow 的充分理由。除了性能方面的優勢，TensorFlow Serving 的另一個重要特性是無需中斷服務，就能實現模型的熱插拔。Zendesk 的這篇博客文章介紹了使用 TensorFlow Serving 部署一個問答機器人的案例：https://medium.com/zendesk-engineering/how-zendesk-serves-tensorflow-models-in-production-751ee22f0f4b

文檔

獲勝者：平局

對於這兩種框架，我都找到了我需要的一切。Python API 的文檔做得很好，兩個框架也都有足夠多的示例和教程可以學習。

一個不太重要的麻煩是 PyTorch 的 C 語言庫基本上沒有文檔。但是，只有當你編寫自定義 C 語言擴展或為這個軟體庫貢獻代碼時才有用。

數據載入

獲勝者：PyTorch

PyTorch 的數據載入 API 設計得很好。數據集、採樣器和數據載入器的介面都是特定的。數據載入器可以接收一個數據集和一個採樣器，並根據該採樣器的調度得出數據集上的一個迭代器（iterator）。並行化數據載入很簡單，只需為數據載入器傳遞一個 num_workers 參數即可。

我還沒找到 TensorFlow 的非常有用的數據載入工具（讀取器、隊列、隊列運行器等等）。部分原因是要將你想並行運行的所有預處理代碼加入到 TensorFlow 圖中並不總是那麼簡單直接（比如計算頻譜圖）。另外，TensorFlow 的 API 本身也更加冗長，學習起來也更難。

設備管理

獲勝者：TensorFlow

TensorFlow 的設備管理的無縫性能非常好。通常你不需要指定任何東西，因為默認的設置就很好。比如說，TensorFlow 假設如果存在可用的 GPU，你就希望在 GPU 上運行。而在 PyTorch 中，你必須在啟用了 CUDA 之後明確地將所有東西移到 GPU 上。

TensorFlow 設備管理的唯一缺陷是它會默認佔用所有可用的 GPU 上的所有內存，即使真正用到的只有其中一個。但也有一種簡單的解決方案，就是指定 CUDA_VISIBLE_DEVICES。有時候人們會忘記這一點，就會讓 GPU 看起來很繁忙，儘管實際上它們啥也沒幹。

在使用 PyTorch 時，我發現我的代碼需要更頻繁地檢查 CUDA 的可用性和更明確的設備管理。尤其是當編寫可以在 CPU 和 GPU 上同時運行的代碼時更是如此。另外，要將 GPU 上的 PyTorch Variable 等轉換成 NumPy 數組也較為繁瑣。

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

自定義擴展

獲勝者：PyTorch

這兩種框架都可以構建或綁定用 C、C++ 或 CUDA 寫的擴展。TensorFlow 還是需要更多樣板代碼，儘管有人認為它能更簡單清晰地支持多種類型和設備。在 PyTorch 中，你只需要簡單地為每個 CPU 和 GPU 版本寫一個介面和對應實現即可。這兩種框架對擴展的編譯都很直接，不需要下載 pip 安裝之外的任何頭文件或源代碼。

關於 TensorBoard 的一點說明

TensorBoard 是一個用於可視化訓練機器學習模型各個方面的工具。它是 TensorFlow 項目產出的最有用的功能之一。僅需在訓練腳本中加入少許代碼，你就可以查看任何模型的訓練曲線和驗證結果。TensorBoard 作為一個網頁服務運行，可以尤其方便地可視化存儲在 headless 節點上的結果。

這是我在使用 PyTorch 時也想繼續使用的一個功能（或找到可替代的工具）。幸運的是，確實有這樣的工具——至少有兩個開源工具可以做到：

• tensorboard_logger：https://github.com/TeamHG-Memex/tensorboard_logger

• crayon：https://github.com/torrvision/crayon

tensorboard_logger 庫用起來甚至比 TensorFlow 中的 TensorBoard「summaries」還簡單，儘管你需要在安裝了 TensorBoard 後才能使用它。crayon 項目可以完全替代 TensorBoard，但需要更多設置（docker 是必需的前提）。

關於 Keras 的一點說明

Keras 是一種帶有可配置的後端的更高層的 API。目前支持 TensorFlow、Theano 和 CNTK，儘管也許不久之後 PyTorch 也將加入這一名單。Keras 也作為 tf.contrib 通過 TensorFlow 提供。

儘管前面我沒有討論 Keras，但這個 API 的使用尤其簡單。它是運行許多常用深度神經網路架構的最快方式。話雖如此，這個 API 並沒有 PyTorch 或核心 TensorFlow 那麼靈活。

關於 TensorFlow Fold 的一點說明

2017 年 2 月，谷歌發布了 TensorFlow Fold。這個庫構建於 TensorFlow 之上，允許實現更動態的圖構建。這個庫的主要優勢應該是動態批量化處理（dynamic batching）。動態批量化可以自動批量化處理不同規模的輸入的計算（考慮一下解析樹上的遞歸網路）。從可編程性上看，它的句法並沒有 PyTorch 的那麼簡單，儘管考慮到批量化在一些情況下帶來的性能提升，這樣的成本也是值得的。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！