对比PyTorch、TensorFlow、JAX、Theano，我发现都在关注两大问题

OneFlow深度学习框架

已于 2022-12-21 10:11:10 修改

阅读量5.6k

点赞数 2

分类专栏：前沿技术文章标签： pytorch tensorflow 人工智能 python 深度学习

于 2022-12-20 08:03:22 首次发布

本文链接：https://blog.csdn.net/oneflow_official/article/details/128391980

版权

本文对比了PyTorch、TensorFlow、JAX和Theano在自动求导、并行处理以及异构计算上的实现，强调了函数转换（如vmap、pmap、pjit）和动态控制流的重要性。各框架通过不同方式处理控制流，例如TensorFlow的tf.cond和JAX的jax.lax.cond。此外，文章讨论了分布式计算的Python-native和分布式运行时两种方法，以及它们在PyTorch和TensorFlow中的应用。

摘要由CSDN通过智能技术生成

作者｜王益

OneFlow社区编译

翻译｜杨婷

最近，我在处理 PyTorch 分布式和 TorchRec 相关的工作，为此，我开始学习 PyTorch 2.0。在业余时间，我也在跟着Alpa作者学习JAX和XLA。如今回顾这些技术，我发现它们的关注点似乎都是如下两个问题：

包含自动求导和并行在内的函数转换，例如 vmap, pmap 和 pjit 等；
异构计算，CPU 负责控制流，GPU/TPU 负责张量计算和集合通信。

本文档中的所有例子都支持在 Colab 中运行：

Theano/Aesara	https://colab.research.google.com/drive/1eg7C5WMNokhXgXQ46pNA30dXUCklquPz
TensorFlow 1.x	https://colab.research.google.com/drive/1jc0ePg2AAXBihevtoZM_33mmhC70rzqz?usp=sharing
TensorFlow 2.x	https://colab.research.google.com/drive/1PbftzJ9E2_FyIiuozTpExMvlFky_G2nv
PyTorch 1.x	https://colab.research.google.com/drive/1v4hENL-IJ-C6VT5H9W1NC2te85D8VdJK
JAX	https://colab.research.google.com/drive/1PlFijLIzAttIBd3tBjiEbSgPXvq9lVlg
functorch/PyTorch 2.x	https://colab.research.google.com/drive/1o-yJ-5g1V084RDaiRw2PqfAjOG7Ty951

函数转换

“函数转换”意为将一个程序转变成另一个程序，最常见的例子是自动求导（autograd）。自动求导采用用户编写的前向过程并创建后向过程，对于用户来说，编写自动求导通常都太过复杂。函数转换的主要难点在于：在编写函数转换算法时以何种方式表示输入和输出过程。

Theano：显式地构建 IR

Theano是最早的深度学习工具之一，也就是如今为人们所熟知的Aesara项目。Theano有一个允许用户在内存中将IR构建为数据结构的API，因此Theano可实现自动求导，并将结果输出为 Python 函数。

import aesara
from aesara import tensor as at


a = at.dscalar("a") # Define placeholders, which have no values.
b = at.dscalar("b")


c = a * b              # c now contains the IR of an expression.TT
dc = aesara.grad(c, a) # Convert the IR in c into another one, dc


f_dc = aesara.function([a, b], dc) # Convert the IR into a Python function,
assert f_dc(1.5, 2.5) == 2.5       # so we can call it.

TensorFlow 1.x：用于运行 IR 的虚拟机

TensorFlow 1.x明确保留了构建IR的想法。若在TensorFlow中运行上述示例，结果不会有什么差别；但倘若在TensorFlow 1.x中来运行，最大的差别在于：