Keras 3迁移指南：从Keras 2到多后端Keras 3的平滑过渡

Keras 3迁移指南：从Keras 2到多后端Keras 3的平滑过渡

keras-io Keras documentation, hosted live at keras.io 项目地址: https://gitcode.com/gh_mirrors/ke/keras-io

前言

随着深度学习框架的不断发展，Keras 3带来了重大更新，支持多后端运行（TensorFlow、JAX和PyTorch）。本文将详细介绍如何将现有的Keras 2代码迁移到Keras 3，帮助开发者充分利用新版本的优势。

环境准备

首先需要安装Keras 3的nightly版本：

pip install keras-nightly

然后设置后端环境：

import os
os.environ["KERAS_BACKEND"] = "tensorflow"  # 可替换为"jax"或"torch"
import keras

基础迁移步骤

1. 导入语句变更：

   • 将from tensorflow import keras改为import keras
   • 将from tensorflow.keras import layers改为from keras import layers
   • 将tf.keras.*改为keras.*

2. 测试现有代码：大多数情况下代码可以直接运行

常见问题及解决方案

1. GPU默认启用JIT编译

Keras 3在GPU上默认启用JIT编译，可能导致某些TensorFlow操作不支持。

解决方案：

model.compile(..., jit_compile=False)
# 或
model.jit_compile = False

2. SavedModel格式变更

Keras 3不再支持直接保存为TF SavedModel格式。

解决方案：

model.export("saved_model")  # 替代model.save()

3. 加载SavedModel

不再支持keras.models.load_model()加载SavedModel。

解决方案：

keras.layers.TFSMLayer("saved_model", call_endpoint="serving_default")

4. 函数式模型输入限制

函数式模型不再支持深度嵌套输入（超过1层）。

解决方案：

# 错误示例
inputs = {
    "foo": keras.Input(...),
    "bar": {"baz": keras.Input(...)}  # 嵌套过深
}

# 正确示例
inputs = {
    "foo": keras.Input(...),
    "bar": keras.Input(...)  # 仅1层嵌套
}

5. TF自动图功能变更

Keras 3不再默认启用TF autograph。

解决方案：

class MyLayer(keras.layers.Layer):
    @tf.function()  # 添加装饰器
    def call(self, inputs):
        ...

6. KerasTensor操作限制

不能在模型构建时直接使用TF操作处理KerasTensor。

解决方案：

# 错误示例
input = keras.layers.Input(...)
tf.squeeze(input)

# 正确示例
input = keras.layers.Input(...)
keras.ops.squeeze(input)

7. 多输出模型评估

Keras 3不再自动返回各输出损失。

解决方案：

model.compile(...,
    metrics=["categorical_crossentropy", "categorical_crossentropy"])

8. 变量跟踪机制变更

直接设置tf.Variable不会被自动跟踪。

解决方案：

# 使用add_weight方法
self.w = self.add_weight(shape=..., initializer="zeros")
# 或使用keras.Variable
self.w = keras.Variable(initial_value=...)

9. None值处理

不允许在嵌套参数中使用None值。

解决方案：

# 方案1：替换None值
inputs = {
    "foo": keras.Input(...),
    "bar": {"baz": keras.Input(...)}  # 不为None
}

# 方案2：使用可选参数
def call(self, foo, baz=None):
    ...

10. 状态构建时机

Keras 3严格要求状态构建时机。

解决方案：

class MyLayer(keras.layers.Layer):
    def build(self, input_shape):  # 在此构建状态
        self.w = self.add_weight(...)
        ...

迁移到多后端

完成上述迁移后，只需更改后端环境变量即可切换到JAX或PyTorch：

os.environ["KERAS_BACKEND"] = "jax"  # 或"torch"

结语

Keras 3的迁移过程相对平滑，大多数情况下只需少量修改。本文涵盖了常见问题和解决方案，帮助开发者顺利完成迁移。迁移后，您将能够充分利用Keras 3的多后端支持，在不同深度学习框架间灵活切换。

keras-io Keras documentation, hosted live at keras.io 项目地址: https://gitcode.com/gh_mirrors/ke/keras-io