TensorFlow VS. PyTorch

本文来源自斯坦福大学计算机科学博士生 Awni Hannun 发表的文章,对比当前两个深度学习主流框架:TensorFlowPytorch 之间的差异。整体而言:

PyTorch 更适合于在研究中快速进行原型设计、业余爱好者和小型项目;而 TensorFlow 则更适合大规模的调度,尤其当考虑到跨平台和嵌入式调度操作时。细分对比如下 ——

准备时间 –> PyTorch

PyTorch 实际上是 NumPy 的替代,它支持 GPU,有着更高级的功能,可以用来构建和训练深度神经网络。如果你熟悉 NumPy、Python 和常见的深度学习概念(卷积层、递归层、SGD 等),那么学习 PyTorch 对你来说会很容易。

而 TensorFlow 则可以看成是一种嵌入 Python 的编程语言。当你在编写 TensorFlow 代码时,它们会通过 Python 编译成一张图,然后由 TensorFlow 执行引擎运行。我看到过刚使用 TensorFlow 的人因为这个额外的间接层而苦思冥想。也因为如此,想用 TensorFlow 还需要学一些额外的概念,比如会话、图、变量作用域和占位符等。要运行基本模型,还需要更多的样板代码。使用 TensorFlow 的前期准备时间肯定比 PyTorch 要长。

图创建和调试 –> PyTorch

创建和运行计算图可能是这两个框架最大的不同。在 PyTorch 中,图架构是动态的,这意味着图是在运行时创建的。而在 TensorFlow 中,图架构是静态的,这意味着先编译出图然后再运行。

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 的 contrib 包中有更多比 PyTorch 更高级的函数和模型。

序列化 –> TensorFlow

在这两种框架下保存和加载模型都很简单。

PyTorch 有一个特别简单的 API,既可以保存模型的所有权重,也可以 pickle 全部类。

TensorFlow 的 Saver 对象也很容易使用,并为检查点(check-pointing)提供了更多选择。在序列化中 TensorFlow 的主要优点是可以将整个图保存为协议缓冲区。这包括参数和运算。此外,该图可以通过其他支持的语言(C++,Java)加载。这对不支持 Python 的调度栈来说至关重要。理论上,在改变模型源代码之后,你想要运行旧模型时它也能有所帮助。

部署 –> TensorFlow

对于小规模的服务器端部署,两个框架都很容易封装在诸如 Flask web 服务器中。不过,TensorFlow 支持移动和嵌入式部署。可以确定的说,这比包括 PyTorch 在内的大多数深度学习框架支持功能的都要多。将 TensorFlow 部署到 Android 或 iOS 上确实需要大量的工作,但至少你不必用 Java 或 C++ 重写模型的整个推理程序。

文档 –> SAME

两种框架的文档中都能找到需要的东西。Python 的 API 在两个框架中都有良好的文档记录,并且有足够的例子和教程来学习这两种框架。一个比较边缘的问题是,PyTorch 的 C 语言库大多是无文档记录的,不过这只影响到编写定制的 C 语言扩展程序,而且这种操作是否有助于软件还存疑。

数据加载 –> PyTorch

PyTorch 中用于数据加载的 API 设计得很好。接口在数据集、采样器和数据加载器中有明确规定。数据加载器接收数据集和采样器,根据采样器的调度,在数据集上生成迭代器。加载并行数据就像把  num_workers 语句传递到数据加载器一样简单。

在 TensorFlow 还没发现特别有用的加载数据的工具,例如 readers, queues, queue runners 等,都不够好。部分原因是因为将想要运行的所有预处理代码添加到 TensorFlow 图中并不总是直接的,例如计算时频谱(spectrogram)。而且,API 本身更繁琐,更难以学习。

设备管理 –> TensorFlow

TensorFlow 管理设备时的无缝性非常好。通常不需要规定任何东西,因为默认已经设好了。例如,如果 GPU 可用,TensorFlow 将默认在 GPU 上运行。在 PyTorch 中,即使支持 CUDA,都必须明确地将所有东西移到设备上。

TensorFlow 设备管理的唯一缺点是,即使你只使用一个 GPU 它也会默认占用所有 GPU 的显存。简单的解决方法是用 CUDA_VISIBLE_DEVICES 语句指定显卡。但有时会忘了设置,所以当 GPU 实际上处于空闲状态时,会显示内存不足。

在 PyTorch 中,代码需要更频繁地检查 CUDA 的可用性和更明确的设备管理,当编写能够同时在 CPU 和 GPU 上运行的代码时尤甚。另外,将 GPU 上的 PyTorch Variable 转换为 NumPy 数组有点繁琐。

自定义扩展 –> PyTorch

在这两种框架中都可以用 C 语言、C++ 或 CUDA 构建或绑定自定义扩展。但 TensorFlow 需要更多的样板代码,即使它支持多种类型和设备。

在 PyTorch 中,只需为每个 CPU 和 GPU 版本编写一个接口和相应的实现。用这两种框架来编译扩展都很直接,并且不需要下载除了 pip 安装包之外的任何头文件或源代码。

TensorFlow 与 PyTorch 入门课件

最后,香港科技大学发布了 TensorFlow 与 Pytorch 两个框架的快速入门课件,以供科研学习者参考。

GitHub地址项目GitHub地址