Keras 3.0预览版迎来重大更新：适用于TensorFlow、JAX和PyTorch

机器之心报道

编辑：陈萍

以后，Keras 适用的框架只会越来越多。

Keras 出新库了，这次是 Keras Core，我们可以将其理解为 Keras 3.0 预览版，预计今年秋天正式发布。

总结而言，这次 Keras Core 对 Keras 代码库进行了完全重写，除了 TensorFlow 之外，它还将 Keras API 引入 JAX 和 PyTorch。

现在，Keras Core 可以作为 tf.keras 的替代品，当使用 TensorFlow 后端时，几乎完全向后兼容 tf.keras 代码。在绝大多数情况下，你只需使用 import keras_core as keras 来代替 from tensorflow import keras 导入即可，替换之后，以前的代码还能正常跑，不仅如此，性能还提高了。

其实，在此之前，Keras 就已经能在 Theano、CNTK（甚至 MXNet）等框架之上运行。最近几年，随着使用 TensorFlow、PyTorch、JAX 的用户越来越多。例如根据 2023 年 StackOverflow 以及 2022 年的 Kaggle 调查数据显示，2022-2023 年 TensorFlow 占据 55% 到 60% 的市场份额，PyTorch 占据 40% 到 45%。同时，JAX 虽然市场份额较小，但已经被谷歌 DeepMind、Midjourney、Cohere 等顶级生成式人工智能公司所接受和采用。

这么看来，本次新库的发布，是在情理之中。

Keras 创始人 François Chollet 表示：现在，（借助 Keras Core）你可以编写跨框架深度学习组件，并从每个框架提供的最佳功能中受益。

接下来，我们看看 Keras Core 具体有哪些特点。

Keras Core 的主要特点

首先是 Keras Core 实现了完整的 Keras API，可支持 TensorFlow、JAX 和 PyTorch。

Keras Core 的第二个特点：它是一个可用于深度学习的、跨框架的低级语言。基于 Keras Core，用户可以创建组件（例如自定义层和预训练模型），而这些组件可以适用于任何框架。特别是，Keras Core 允许用户访问 keras_core.ops 命名空间，后者适用于所有后端。

Keras Core 的第三个特点是与 JAX、PyTorch 和 TensorFlow 中的本地工作流无缝集成。与 Keras 1.0 不同，Keras Core 不只是针对以 Keras 为中心的工作流，还意味着可以与低级的后端本地工作流无缝地工作。

Keras Core 其他特点还包括：支持所有后端跨框架的数据 pipeline。多框架的机器学习意味着多框架的数据加载和处理，处理起来比较麻烦。现在 Keras Core 模型可以使用广泛的数据 pipeline 进行训练 —— 不管你是使用 JAX、PyTorch 还是 TensorFlow 后端。

预训练模型。从现在开始，你可以借助 Keras Core 使用更多的预训练模型。现在已经有 40 个 Keras 应用模型可在后端中使用，此外，KerasCV 和 KerasNLP 中存在的大量预训练模型（例如 BERT、T5、YOLOv8、Whisper 等）也适用于所有后端。