Keras Core: 下一代多后端深度学习框架

Ray

Keras Core:开启深度学习新时代

Keras长期以来一直是深度学习领域最受欢迎的高级API之一。如今,随着Keras Core的问世,这一强大工具迎来了新的进化。Keras Core是Keras API的多后端实现,支持TensorFlow、JAX和PyTorch,为深度学习开发者提供了前所未有的灵活性和可移植性。

什么是Keras Core?

Keras Core是一个全新的多后端Keras实现,允许开发者在TensorFlow、JAX和PyTorch之间无缝切换。它保留了Keras简洁易用的API风格,同时赋予用户选择最适合自己需求的后端框架的能力。无论是追求JAX的高性能,还是需要TensorFlow完善的生产部署工具,亦或是偏爱PyTorch的动态计算图,Keras Core都能满足您的需求。

Keras Core架构

Keras Core的主要特性

多后端支持: 可以在TensorFlow、JAX和PyTorch之间自由切换,充分利用各个框架的优势。
统一的API: 保持了Keras简洁直观的API设计,降低学习成本。
高度兼容性: 大多数现有的Keras模型可以直接在Keras Core上运行,实现平滑迁移。
跨框架组件: 可以创建在任何支持的后端上都能工作的自定义层、模型和指标。
灵活的数据管道: 支持各种数据格式,包括tf.data、PyTorch DataLoader等。
分布式训练: 提供新的分布式API,支持数据并行和模型并行。

为什么选择Keras Core?

性能优化: 可以根据具体模型选择性能最佳的后端,无需修改代码。
生态系统灵活性: 可以利用不同框架的生态系统优势,如PyTorch的扩展包、TensorFlow的部署工具等。
大规模训练: 利用JAX进行高效的模型并行和数据并行分布式训练。
代码可移植性: 一次编写,可在多个框架上运行,提高代码复用性。
未来兼容: 避免被单一框架锁定,保持技术选择的灵活性。

Keras Core的工作原理

Keras Core通过抽象出底层计算逻辑,提供了一个统一的高级API。它包含了一个名为keras.ops的命名空间,实现了完整的NumPy API以及一些神经网络特定的操作。只要使用这些ops,您的自定义组件就可以在所有支持的后端上无缝运行。

import os
os.environ["KERAS_BACKEND"] = "jax"  # 或 "tensorflow" 或 "torch"
import keras_core as keras

# 定义模型
model = keras.Sequential([
    keras.layers.Dense(64, activation="relu"),
    keras.layers.Dense(10, activation="softmax")
])

# 编译模型
model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])

# 训练模型
model.fit(x_train, y_train, epochs=5, validation_split=0.2)

这段代码可以在任何支持的后端上运行,只需更改KERAS_BACKEND环境变量即可切换。