Keras 3: 多框架深度学习API的革新

Keras 3简介

Keras是一个高级神经网络API,专注于实现快速实验。Keras 3是对Keras的全面重写,旨在让用户能够在JAX、TensorFlow或PyTorch等多个后端上运行Keras工作流程,并解锁全新的大规模模型训练和部署能力。

Keras 3保持了Keras一贯的简单易用特性,同时大大增强了其灵活性和功能性。它允许用户根据具体需求选择最合适的后端框架,并在不同后端之间自如切换。此外,Keras 3还可以作为一种低级跨框架语言,用于开发可在JAX、TensorFlow或PyTorch原生工作流中使用的自定义组件。

Keras 3的主要优势

1. 性能优化

Keras 3让用户无需修改代码即可为模型选择最佳性能的后端。根据基准测试,JAX通常在GPU、TPU和CPU上提供最佳的训练和推理性能,但具体结果因模型而异。这种动态选择后端的能力确保了用户能以最高效率进行训练和服务。

2. 生态系统兼容性

任何Keras 3模型都可以实例化为PyTorch Module,导出为TensorFlow SavedModel,或实例化为无状态JAX函数。这意味着用户可以利用PyTorch生态系统包、TensorFlow全套部署和生产工具,以及JAX大规模TPU训练基础设施。用户只需编写一个使用Keras 3 API的model.py,就能访问机器学习世界的所有资源。

3. 大规模并行处理

Keras 3引入了全新的分布式API keras.distribution,目前已在JAX后端实现(TensorFlow和PyTorch后端即将推出)。它使模型并行、数据并行及两者的组合变得简单,可应用于任意规模的模型和集群。由于该API将模型定义、训练逻辑和分片配置完全分离,因此使分布式工作流程易于开发和维护。

4. 开源模型发布效果最大化

使用Keras 3实现的预训练模型可以立即被任何人使用,无论他们选择哪种框架(即使他们不是Keras用户)。这意味着开发者可以用相同的开发成本获得两倍的影响力。

5. 多源数据管道兼容

Keras 3的fit()/evaluate()/predict()例程与tf.data.Dataset对象、PyTorch DataLoader对象、NumPy数组、Pandas数据框兼容,不受使用后端的限制。这意味着用户可以在PyTorch DataLoader上训练Keras 3 + TensorFlow模型,或在tf.data.Dataset上训练Keras 3 + PyTorch模型。

Keras 3的核心功能

全面的Keras API支持

Keras 3实现了完整的Keras API,包括100多个层、数十个指标、损失函数、优化器和回调,以及Keras训练和评估循环、保存和序列化基础设施。所有这些API都可在TensorFlow、JAX和PyTorch上使用。

任何只使用内置层的Keras模型都可以立即在所有支持的后端上工作。事实上,只使用内置层的现有tf.keras模型可以立即在JAX和PyTorch上运行。

多框架组件开发

Keras 3使用户能够创建在任何框架中以相同方式工作的组件(如任意自定义层或预训练模型)。特别是,Keras 3提供了适用于所有后端的keras.ops命名空间,其中包含:

• NumPy API的完整实现
• 神经网络特定函数,如ops.softmax、ops.binary_crossentropy、ops.conv等

只要用户仅使用keras.ops中的操作,他们的自定义层、自定义损失、自定义指标和自定义优化器就可以在JAX、PyTorch和TensorFlow中使用相同的代码工作,并产生完全相同的数值结果。

与原生工作流的无缝集成

Keras 3不仅适用于Keras中心化工作流,还可以无缝集成到低级后端原生工作流中。用户可以将Keras模型(或任何其他组件,如损失或指标)用于JAX训练循环、TensorFlow训练循环或PyTorch训练循环,或作为JAX或PyTorch模型的一部分,几乎没有任何阻力。

例如,用户可以:

• 编写低级JAX训练循环,使用optax优化器、jax.grad、jax.jit、jax.pmap训练Keras模型
• 编写低级TensorFlow训练循环,使用tf.GradientTape和tf.distribute训练Keras模型
• 编写低级PyTorch训练循环,使用torch.optim优化器、torch损失函数和torch.nn.parallel.DistributedDataParallel包装器训练Keras模型
• 在PyTorch Module中使用Keras层(因为它们也是Module实例)
• 在Keras模型中像使用Keras层一样使用任何PyTorch Module

预训练模型支持

Keras 3为用户提供了广泛的预训练模型选择。所有40个Keras Applications模型(keras.applications命名空间)都可在所有后端使用。此外,KerasCV和KerasNLP中的大量预训练模型也可与所有后端兼容,包括:

• BERT
• OPT
• Whisper
• T5
• StableDiffusion
• YOLOv8
• SegmentAnything 等

总结

Keras 3代表了深度学习API的一次重大革新。通过支持多个后端框架,它为用户提供了前所未有的灵活性和功能性。无论是对于快速原型设计、大规模模型训练还是生产部署,Keras 3都提供了强大的工具和简洁的API。它不仅保持了Keras一贯的易用性,还大大拓展了其应用范围,使其成为一个真正的多框架深度学习解决方案。

随着深度学习技术的不断发展,Keras 3将继续发挥其作为连接不同框架和生态系统的桥梁作用,为研究人员和工程师提供更多可能性。无论您是深度学习新手还是经验丰富的专家,Keras 3都将成为您强大而灵活的开发工具。