论文 TENSORFLOW.JS: MACHINE LEARNING FOR THE WEB AND BEYOND 指出:
由于我们设计目标的一个重要部分是优先考虑易用性而不是性能,因此TensorFlow.js支持急切风格的微分运算。
一般来说,我们在讨论什么样的性能损失?这是否取决于模型?是否有没有任何性能差异的情况?
回答:
声明式(也称为延迟执行、图模式)编程范式(如TensorFlow v1(Python)默认图模型中的那种)的主要性能优势来自以下几个方面:
- 将整个模型推送到C++层,在那里执行开销与解释型或非编译型语言(如Python和JavaScript)相比要低得多
- 并行执行模型计算图中独立路径的一个例子是一个由多个独立输入塔组成的模型。这些塔可以在CPU的不同核心或同一主机的多个GPU上并发执行。
- 由于在执行开始之前就已知整个模型,C++执行引擎可以对模型的计算图进行一整套优化。举几个例子:
- 常量折叠:由仅对常量节点进行无状态、确定性操作的子树可以折叠成单个常量节点
- 操作融合:在某些情况下,计算图中的几个相邻节点(操作)可以被数学上等价但计算上更有效的节点所替代。
- 修剪:某些计算图包含不对最终输出做出贡献的节点。图模型执行引擎可以预先看到这一点并防止这些节点执行。
- 即时(JIT)编译:图执行引擎可以获取整个图并将其编译为低级表示,这种表示涉及较低的调度开销,并且更适合在可用硬件上进行高性能执行(例如,针对NVIDIA或兼容GPU的CUDA程序,针对Google TPU的特殊指令,甚至是针对WebGL的着色器程序等)
所有上述优化都由图模式TensorFlow支持。有关更多详细信息,请谷歌搜索“grappler”和“XLA”。
TensorFlow.js采用了命令式(也称为急切)范式,主要基于可用性考虑。这与TensorFlow急切执行、PyTorch和NumPy类似。因此,它不具备上述所有优化机会。
然而,请注意,从命令式程序中获取计算图的方法是存在的(参见TensorFlow v2的tf.function装饰器和 JAX)。TensorFlow.js没有理由不能采用类似的范式来提升性能。只是产品团队尚未明确需要优先考虑这一功能。