tensorflow - 简单来说，JAX、Trax 和 TensorRT 有什么区别？

Question

我一直在使用 TensorRT 和 TensorFlow-TRT 来加速我的 DL 算法的推理。

然后我听说了：

• JAX https://github.com/google/jax
• 特拉克斯https://github.com/google/trax

两者似乎都加速了深度学习。但我很难理解他们。任何人都可以简单地解释它们吗？

Answer 1

Trax是由谷歌创建并被谷歌大脑团队广泛使用的深度学习框架。TensorFlow原则上，PyTorch它是自然语言处理领域的现有技术深度学习模型（例如 Transformers、Bert 等）的替代方案。

Trax建立在TensorFlow和之上JAX。JAX是 Numpy 的增强和优化版本。JAX和的重要区别NumPy在于前者使用了一个名为 XLA（高级线性代数）的库，它允许在上面运行NumPy代码GPU，TPU而不是CPU像在 plain 中那样运行代码NumPy，从而加快计算速度。