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

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

作者說：

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

一個(gè)具體的例子是，在 PyTorch 和 TensorFlow 兩個(gè)庫中的 Adam 實(shí)現(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 有一個(gè)專用的 C++ 實(shí)現(xiàn)。因此，在這里，TensorFlow 不會(huì)在 Python 上花費(fèi)額外的時(shí)間，而且它在 C++ 中有一個(gè)優(yōu)化的實(shí)現(xiàn)。在這種情況下，為什么 TensorFlow 版本沒有更快一些呢？

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

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

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

SR2Z：

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

huberloss (回復(fù) SR2Z)：

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

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

entarko：

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

patrickkidger：

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

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

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

programmerChilli (回復(fù) patrickkidger)：

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

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

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

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

un_om_de_cal：

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

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

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