Skip to main content
 首页 » 编程设计

python之pytorch 能否优化顺序操作(如 tensorflow 图或 JAX 的 jit)

2024年09月03日33exmyth

本来tensorflow和pytorch有一个根本的区别:

  • tensorflow 基于计算图。构建此图和在 session 中对其进行评估是两个独立的步骤。在使用过程中,图形不会改变,因此可以进行优化。
  • torch 热切地评估张量上的操作。这使 API 更方便(无 session ),但也失去了识别和优化始终按顺序发生的操作的潜力。

现在这种差异变得不那么明显了。 Tensorflow 用 tf eager 回应了 torch 的流行.还有 JAX项目,它建立在与 tensorflow ( XLA) 相同的底层框架上。 JAX 没有 session 的概念。但它允许您通过简单地调用 jit 将多个操作编译在一起。 .

既然 Tensorflow 已经覆盖 PyTorch 功能,那么 PyTorch 是否也在努力整合 Tensorflow 的优势? PyTorch(或其路线图)中是否有类似 session 或 jit 功能的东西?

API 文档有一个 jit section ,但据我所知,这更多是关于导出模型。

请您参考如下方法:

如您所述,有一个 torch.jit它的目的也是在导出的图中引入优化(例如内核融合、常量优化等)。 IIRC 你可以在他们的 github repo here 中找到一些源代码。 ,尽管我不确定文档中的某处是否明确提到了这些内容(或明确到足以让人记住)。

从 1.3 开始,还引入了量化(有关介绍,请参阅 here)。在教程部分,即here您可以看到 Conv2d、BatchNorm 和 ReLU 的显式融合以提高性能。 Ofc 也存在特定的东西,比如使用 int 而不是 float 来进行权重(量化),混合算术(尽可能使用 half float 精度,见NVidia 的 Apex ) 等。

最后但同样重要的是,我认为对于使用矢量化操作并使用 torchscript 导出的编写良好的模型,您不会看到真正显着的运行时差异因为一些通用图形优化。仍然不同的是你是否要使用 GPU、CPU、TPU,它们的版本是什么,你是仅在推理还是训练之后等等。很难确定 tensorflow 比较的速度有多快到 pytorch(除了两个框架中的一些众所周知的问题)。总而言之,这取决于 AFAIK 测量结果差异很大。

顺便说一句。当谈到每个框架的优势时,它们的核心确实开始涵盖类似的东西(PyTorch 最近获得了移动支持,请参阅 here)。真正的区别仍然是不同的底层方法以及每个框架必须做什么来规避这些限制。