PyTorch可以和TensorFlow一样快，有时甚至比TensorFlow更快了？

PyTorch可以和TensorFlow一样快，有时甚至比TensorFlow更快了？这是怎么回事？最近Reddit的一个帖子引起热议。

近日，Reddit 上有一个热帖：为什么 PyTorch 和 TensorFlow 一样快 (有时甚至比 TensorFlow 更快)？

作者说：

由于这两个库都在底层使用 cuDNN，所以我预想各个操作的速度是相似的。然而，TensorFlow (在 graph 模式下) 编译一个 graph，因此当你运行实际的训练循环时，在 session.run 调用之外没有任何 Python 开销。在 PyTorch 中，由于动态图的关系，需要经常使用 Python，所以我预想这回增加一些开销。更不用说，拥有静态图意味着可以进行图优化，比如节点修剪和排序操作。但我在网上看到的许多基准测试中，在 GPU 上，PyTorch 都可以轻松地赶上 TensorFlow。

一个具体的例子是，在 PyTorch 和 TensorFlow 两个库中的 Adam 实现：

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 (据我所知这是常见情况), TensorFlow 有一个专用的 C++ 实现。因此，在这里，TensorFlow 不会在 Python 上花费额外的时间，而且它在 C++ 中有一个优化的实现。在这种情况下，为什么 TensorFlow 版本没有更快一些呢？

我听说 PyTorch 在 cuDNN 级别上进行了更好的优化。有人能提供更多细节吗？是什么阻止了 TensorFlow 做同样的事情？我所知道的惟一优化是 PyTorch 使用 NCHW 格式 (针对 cuDNN 进行了更好的优化)，而 TensorFlow 默认使用 NHWC。

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

这个问题引发了诸多 Reddit 网友的关注，大家纷纷给出自己的答案。

SR2Z：

这里的关键是异步执行 —— 除非你不断地在 GPU 之间复制数据，否则 PyTorch 操作只会为 GPU 排队。对 torch 函数的 Python 调用将在排队操作后返回，因此大多数 GPU 工作都不会占用 Python 代码。这将瓶颈从 Python 转移到了 CUDA，这就是为什么它们执行起来如此相似。

huberloss (回复 SR2Z)：

TF 构建了一个执行图，然后由 C++ 后端处理，这意味着你在 Python 中所做的唯一工作就是设置这个图 (至少在 TF1.x 中是这样)。我不认为 TF 会移动移入 / 移出 GPU，除非需要 (例如，op1 运行在 CPU 上，op2 运行在 GPU 上，op3 运行在 CPU 上 -> 这将导致向 GPU 复制或从 GPU 复制）。TF 有执行程序，它也将异步运行所有内容 (这意味着独立节点将以一种你能够尽快满足输出节点的方式执行)。

关于 OP 的观点，我真的不知道为什么有时候会更快。我可以想到的唯一猜测是数据格式，或者某些 ops 调用 CUDA/cuDNN 的方式。

entarko：

正如你所说，这两个库都使用 cuDNN，所以在较低级别上使用的大多数算法是相似的。我的理解是，在默认情况下，PyTorch 中的执行是异步的。这意味着 Python 命令在 GPU 上执行某项调用，但不等待该调用的结果，除非下一个操作需要该结果。因此，GPU 操作速度的唯一不同之处在于 python 调用所需的时间，这与 GPU 上的实际计算相比总体上较小。这并不能解释为什么有时候 PyTorch 更快，这在一定程度上来自 NCHW (尽管并不总是这样)，而且还来自算法为不同的操作进行的选择 (cuDNN 的 benchmark flag)。

patrickkidger：

这与 PyTorch 和 TensorFlow 没有直接关系，但是既然 NCHW 和 NHWC 被提了出来，我希望这里有人能知道答案……

为什么 batch size N 是第一个而不是最后一个？在我看来，由于缓存的局部性，最后使用它通常会加快速度。

大多数数据集格式不统一应该不是大问题：相比于在整个计算的其余部分中具有 N-last 格式的好处而言，单个副本转到 N-last 格式（即在将其送入神经网络之前）的开销应该是微不足道的。

programmerChilli (回复 patrickkidger)：

我认为你是正确的，将批处理维度放在最后可以实现一些非常有趣的优化 (想想所有的向量化)。

但是，如果你的 batch size 很小，会发生什么呢？比如 batch size 为 2 (可能 3 会更糟)。突然之间，你可以在批处理维度上进行向量化的数量变得非常少了，并且你已经传播了其余的数据，但没有获得多少收益。

实际上，以前有几个框架使用这种格式，比如来自 Nervana 的 Neon。

我想主要的答案是人们尝试了，但是由于各种原因，一些政治上 / 技术上的原因，我们目前已经把 batch size 作为第一个维度确定下来。

un_om_de_cal：

根据我将项目从 PyTorch 转移到 TF 的经验，虽然 PyTorch 的基本训练操作 (前向传播、反向传播、更新参数) 更快，但是一些不寻常的操作 (包括 argmax 和 slicing)，PyTorch 比 TF 要慢得多。因此，也许 PyTorch 更适合通常基准测试的更常见的操作，而 TF 针对更广泛的操作进行了优化？

我想到的另一点是，PyTorch 教程在 CPU 上做数据增强，而 TF 教程在 GPU 上做数据增强 (至少 1-2 年前我看到的教程是这样)。如果你像这样做项目，那么你将在 CPU 上执行一些计算，只要你没有耗尽 CPU, 就可以提高效率。

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