[论文解读] SymJAX: symbolic CPU/GPU/TPU programming

Randall Balestriero|arXiv (Cornell University)|May 21, 2020

Music and Audio Processing参考文献 5被引用 4

一句话总结

SymJAX 是一个基于 JAX 和 XLA 的符号编程框架，通过内核融合、减少内存操作，支持自动微分，实现了在 CPU、GPU 和 TPU 上的高效、优化的计算图。它提供了一个高层次的、类似 PyTorch 的 API，支持即时编译和广泛的硬件支持，通过图优化和符号计算展示了性能提升。

ABSTRACT

SymJAX is a symbolic programming version of JAX simplifying graph input/output/updates and providing additional functionalities for general machine learning and deep learning applications. From an user perspective SymJAX provides a la Theano experience with fast graph optimization/compilation and broad hardware support, along with Lasagne-like deep learning functionalities.

研究动机与目标

• 为深度学习提供一个符号编程接口，简化图构建、输入/输出处理和参数更新。
• 通过将多个内核融合为单个操作来优化计算图，减少内存传输并提升性能。
• 通过 XLA 支持 CPU、GPU 和 TPU 硬件上的高效即时编译和自动微分。
• 提供一个高层次、直观的 API，类似于 PyTorch，同时保留 JAX 的性能和编译优势。
• 通过内置优化器、损失函数和层组件支持，实现端到端的训练和推理流水线。

提出的方法

• 使用符号张量（通过 JAX 的 XLA 后端）表示计算，支持静态图构建和优化。
• 通过 XLA 实现内核融合，合并操作，最小化中间内存分配并提升执行速度。
• 通过 JAX 的 XLA 编译器集成即时编译（JIT）以加速计算图。
• 通过 SymJAX 的梯度函数支持自动微分，实现在符号表达式上的反向传播。
• 通过 `sj.function` 提供函数式 API，用于编译和执行指定输出与更新的图。
• 通过 `sj.layers` 和 `sj.optimizers` 提供模块化的神经网络组件（如 Conv2D、BatchNorm、ReLU），用于模型构建。

实验结果

研究问题

• RQ1 基于 JAX 和 XLA 构建的符号编程接口能否在支持多种加速器的同时，实现与 PyTorch 相当的高性能和易用性？
• RQ2 在深度学习工作负载中，内核融合在减少内存操作和提升训练与推理速度方面的效果如何？
• RQ3 符号计算与自动微分在多大程度上可以实现用户友好且可组合，以支持复杂的深度学习模型？
• RQ4 XLA 基础的编译与即时编译的集成在 CPU、GPU 和 TPU 上如何提升深度学习模型的性能？
• RQ5 能否在 JAX 上构建一个高层次、类似 PyTorch 的 API，实现极低性能开销，并完整支持优化和微分？

主要发现

• SymJAX 顺利编译并优化了多变量正态分布概率密度函数的计算，实现了延迟求值并正确输出，使用了符号张量操作。
• 在随机梯度下降实验中，SymJAX 正确计算了梯度并更新了参数，损失值在 10 次迭代中从 ~0.00847 降低至 ~0.0065。
• 在 CIFAR-10 数据集上，SymJAX 训练的模型在 100 个周期后达到测试损失 ~0.815 和测试准确率 ~72.1%，展示了经过优化计算的有效训练。
• 该框架支持复杂的模型架构，包括卷积层、批归一化层、池化层和全连接层，具备正确的形状传播和参数追踪能力。
• 通过使用 `sj.function`，实现了带有更新的图的高效编译与执行，如训练循环中每轮记录损失和准确率所示。
• 内核融合与 JIT 编译显著减少了内存操作并提升了性能，CIFAR-10 训练流水线的高效执行已得到验证。