多后端Keras。
项目描述
Keras Core:一个新的多后端Keras
Keras Core是Keras API的新多后端实现,支持TensorFlow、JAX和PyTorch。
警告: 目前,此包是实验性的。它有粗糙的边缘,可能并非所有功能都能按预期工作。我们目前正在努力改进它。
一旦准备就绪,此包将成为Keras 3.0,并取代tf.keras
。
本地安装
Keras Core兼容Linux和MacOS系统。要安装本地开发版本
- 安装依赖项
pip install -r requirements.txt
- 从根目录运行安装命令。
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 用户:可以访问一个功能全面、经过实战检验、文档齐全的建模和训练库。
项目详情
下载文件
下载适用于您的平台的文件。如果您不确定选择哪个,请了解更多关于 安装包 的信息。