Keras Core：一个新的多后端Keras

Keras Core是Keras API的新多后端实现，支持TensorFlow、JAX和PyTorch。

警告： 目前，此包是实验性的。它有粗糙的边缘，可能并非所有功能都能按预期工作。我们目前正在努力改进它。

一旦准备就绪，此包将成为Keras 3.0，并取代tf.keras。

本地安装

Keras Core兼容Linux和MacOS系统。要安装本地开发版本

1. 安装依赖项

pip install -r requirements.txt

2. 从根目录运行安装命令。

python pip_build.py --install

您还应该安装您选择的后端：tensorflow、jax或torch。请注意，使用某些Keras Core功能需要tensorflow：某些预处理层以及tf.data管道。

配置您的后端

您可以导出环境变量KERAS_BACKEND，或者您可以在~/.keras/keras.json中编辑您本地的配置文件来配置您的后端。可用的后端选项有："tensorflow"、"jax"、"torch"。示例

export KERAS_BACKEND="jax"

在Colab中，您可以这样做

import os
os.environ["KERAS_BACKEND"] = "jax"

import keras_core as keras

注意： 在导入keras_core之前必须配置后端，并且导入包之后不能更改后端。

向后兼容性

Keras 核心旨在作为 tf.keras 的直接替代品（当使用 TensorFlow 后端时）。只需将现有的 tf.keras 代码中的 keras 导入更改为 keras_core，确保使用最新的 .keras 格式调用 model.save()，即可完成。

如果您的 tf.keras 模型不包含自定义组件，您可以直接在 JAX 或 PyTorch 上运行它。

如果它包含自定义组件（例如自定义层或自定义 train_step()），通常只需几分钟即可将其转换为与后端无关的实现。

此外，Keras 模型可以以任何格式消费数据集，无论您使用的是哪个后端：您可以使用现有的 tf.data.Dataset 管道或 PyTorch DataLoaders 训练模型。

为什么使用 Keras 核心？

• 在任意框架之上运行您的 Keras 工作流程，从而能够自由地享受每个框架的优势，例如 JAX 的可扩展性和性能或 TensorFlow 的生产生态系统选项。
• 编写自定义组件（例如层、模型、指标），您可以在任何框架的低级工作流程中使用它们。
  • 您可以将 Keras 模型与从头开始编写的 native TF、JAX 或 PyTorch 训练循环一起训练。
  • 您可以将 Keras 模型用作 PyTorch 原生 Module 的一部分，或作为 JAX 原生模型函数的一部分。
• 通过避免框架锁定，使您的 ML 代码面向未来。
• 作为 PyTorch 用户：终于可以访问 Keras 的强大功能和易用性了！
• 作为 JAX 用户：可以访问一个功能全面、经过实战检验、文档齐全的建模和训练库。