PyTorch可以和TensorFlow一樣快，有時甚至比TensorFlow更快了？這是怎么回事？最近Reddit的一個帖子引起熱議。

近日，Reddit 上有一個熱帖：為什么 PyTorch 和 TensorFlow 一樣快 (有時甚至比 TensorFlow 更快)？

作者說：

由于這兩個庫都在底層使用 cuDNN，所以我預想各個操作的速度是相似的。然而，TensorFlow (在 graph 模式下) 編譯一個 graph，因此當你運行實際的訓練循環(huán)時，在 session.run 調(diào)用之外沒有任何 Python 開銷。在 PyTorch 中，由于動態(tài)圖的關(guān)系，需要經(jīng)常使用 Python，所以我預想這回增加一些開銷。更不用說，擁有靜態(tài)圖意味著可以進行圖優(yōu)化，比如節(jié)點修剪和排序操作。但我在網(wǎng)上看到的許多基準測試中，在 GPU 上，PyTorch 都可以輕松地趕上 TensorFlow。

一個具體的例子是，在 PyTorch 和 TensorFlow 兩個庫中的 Adam 實現(xiàn)：

https://github.com/pytorch/pytorch/blob/master/torch/optim/adam.py https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/adam.py

PyTorch 擁有你所期望的所有操作。對于 {_resource} _apply_density 情況下的 TensorFlow (據(jù)我所知這是常見情況), TensorFlow 有一個專用的 C++ 實現(xiàn)。因此，在這里，TensorFlow 不會在 Python 上花費額外的時間，而且它在 C++ 中有一個優(yōu)化的實現(xiàn)。在這種情況下，為什么 TensorFlow 版本沒有更快一些呢？

我聽說 PyTorch 在 cuDNN 級別上進行了更好的優(yōu)化。有人能提供更多細節(jié)嗎？是什么阻止了 TensorFlow 做同樣的事情？我所知道的惟一優(yōu)化是 PyTorch 使用 NCHW 格式 (針對 cuDNN 進行了更好的優(yōu)化)，而 TensorFlow 默認使用 NHWC。

更新：看起來新的 Volta GPU 使用 NHWC 格式 (TensorFlow 的默認格式) 的性能更好：https://devblogs.nvidia.com/tensor-core-ai-performance- ones/

這個問題引發(fā)了諸多 Reddit 網(wǎng)友的關(guān)注，大家紛紛給出自己的答案。

SR2Z：

這里的關(guān)鍵是異步執(zhí)行 —— 除非你不斷地在 GPU 之間復制數(shù)據(jù)，否則 PyTorch 操作只會為 GPU 排隊。對 torch 函數(shù)的 Python 調(diào)用將在排隊操作后返回，因此大多數(shù) GPU 工作都不會占用 Python 代碼。這將瓶頸從 Python 轉(zhuǎn)移到了 CUDA，這就是為什么它們執(zhí)行起來如此相似。

huberloss (回復 SR2Z)：

TF 構(gòu)建了一個執(zhí)行圖，然后由 C++ 后端處理，這意味著你在 Python 中所做的唯一工作就是設置這個圖 (至少在 TF1.x 中是這樣)。我不認為 TF 會移動移入 / 移出 GPU，除非需要 (例如，op1 運行在 CPU 上，op2 運行在 GPU 上，op3 運行在 CPU 上 -> 這將導致向 GPU 復制或從 GPU 復制）。TF 有執(zhí)行程序，它也將異步運行所有內(nèi)容 (這意味著獨立節(jié)點將以一種你能夠盡快滿足輸出節(jié)點的方式執(zhí)行)。

關(guān)于 OP 的觀點，我真的不知道為什么有時候會更快。我可以想到的唯一猜測是數(shù)據(jù)格式，或者某些 ops 調(diào)用 CUDA/cuDNN 的方式。

entarko：

正如你所說，這兩個庫都使用 cuDNN，所以在較低級別上使用的大多數(shù)算法是相似的。我的理解是，在默認情況下，PyTorch 中的執(zhí)行是異步的。這意味著 Python 命令在 GPU 上執(zhí)行某項調(diào)用，但不等待該調(diào)用的結(jié)果，除非下一個操作需要該結(jié)果。因此，GPU 操作速度的唯一不同之處在于 python 調(diào)用所需的時間，這與 GPU 上的實際計算相比總體上較小。這并不能解釋為什么有時候 PyTorch 更快，這在一定程度上來自 NCHW (盡管并不總是這樣)，而且還來自算法為不同的操作進行的選擇 (cuDNN 的 benchmark flag)。

patrickkidger：

這與 PyTorch 和 TensorFlow 沒有直接關(guān)系，但是既然 NCHW 和 NHWC 被提了出來，我希望這里有人能知道答案……

為什么 batch size N 是第一個而不是最后一個？在我看來，由于緩存的局部性，最后使用它通常會加快速度。

大多數(shù)數(shù)據(jù)集格式不統(tǒng)一應該不是大問題：相比于在整個計算的其余部分中具有 N-last 格式的好處而言，單個副本轉(zhuǎn)到 N-last 格式（即在將其送入神經(jīng)網(wǎng)絡之前）的開銷應該是微不足道的。

programmerChilli (回復 patrickkidger)：

我認為你是正確的，將批處理維度放在最后可以實現(xiàn)一些非常有趣的優(yōu)化 (想想所有的向量化)。

但是，如果你的 batch size 很小，會發(fā)生什么呢？比如 batch size 為 2 (可能 3 會更糟)。突然之間，你可以在批處理維度上進行向量化的數(shù)量變得非常少了，并且你已經(jīng)傳播了其余的數(shù)據(jù)，但沒有獲得多少收益。

實際上，以前有幾個框架使用這種格式，比如來自 Nervana 的 Neon。

我想主要的答案是人們嘗試了，但是由于各種原因，一些政治上 / 技術(shù)上的原因，我們目前已經(jīng)把 batch size 作為第一個維度確定下來。

un_om_de_cal：

根據(jù)我將項目從 PyTorch 轉(zhuǎn)移到 TF 的經(jīng)驗，雖然 PyTorch 的基本訓練操作 (前向傳播、反向傳播、更新參數(shù)) 更快，但是一些不尋常的操作 (包括 argmax 和 slicing)，PyTorch 比 TF 要慢得多。因此，也許 PyTorch 更適合通常基準測試的更常見的操作，而 TF 針對更廣泛的操作進行了優(yōu)化？

我想到的另一點是，PyTorch 教程在 CPU 上做數(shù)據(jù)增強，而 TF 教程在 GPU 上做數(shù)據(jù)增強 (至少 1-2 年前我看到的教程是這樣)。如果你像這樣做項目，那么你將在 CPU 上執(zhí)行一些計算，只要你沒有耗盡 CPU, 就可以提高效率。

各位 PyTorcher 和 TF boy，你們怎么看？