► 开发者指南 / 将 Keras 2 代码迁移到多后端 Keras 3

将 Keras 2 代码迁移到多后端 Keras 3

作者: Divyashree Sreepathihalli
创建日期: 2023/10/23
最后修改日期: 2023/10/30
描述: 迁移您的 Keras 2 代码到多后端 Keras 3 的说明和故障排除。

在 Colab 中查看 • GitHub 源代码

本指南将帮助您将仅支持 TensorFlow 的 Keras 2 代码迁移到多后端 Keras 3 代码。迁移的开销很小。一旦您迁移完成，您可以在 JAX、TensorFlow 或 PyTorch 上运行 Keras 工作流。

本指南分为两个部分：

将您的遗留 Keras 2 代码迁移到 Keras 3，该版本运行于 TensorFlow 后端。这通常非常简单，尽管需要注意一些小问题，我们将详细讨论。
进一步将您的 Keras 3 + TensorFlow 代码迁移到多后端 Keras 3，以便它可以在 JAX 和 PyTorch 上运行。

让我们开始吧。

设置

首先，安装 keras-nightly。

此示例使用 TensorFlow 后端（os.environ["KERAS_BACKEND"] = "tensorflow"）。迁移代码后，您可以将 "tensorflow" 字符串更改为 "jax" 或 "torch"，然后在 Colab 中单击 “重新启动运行时”，您的代码将会在 JAX 或 PyTorch 后端上运行。

!pip install -q keras-nightly

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import keras
import tensorflow as tf
import numpy as np

 [notice] pip 有新版本可用：23.3.1 -> 24.0
 [notice] 要更新，请运行：pip install --upgrade pip

从 Keras 2 转到 Keras 3（使用 TensorFlow 后端）

首先，替换您的导入：

将 from tensorflow import keras 替换为 import keras
将 from tensorflow.keras import xyz（例如 from tensorflow.keras import layers）替换为 from keras import xyz（例如 from keras import layers）
将 tf.keras.* 替换为 keras.*

接下来，开始运行您的测试。大多数情况下，您的代码将在 Keras 3 上正常执行。您可能遇到的所有问题已在下面详细列出，并提供了解决方案。

在 GPU 上，`jit_compile` 默认设置为 `True`。

在 Keras 3 中，Model 构造函数的 jit_compile 参数的默认值已在 GPU 上设置为 True。这意味着模型将在 GPU 上默认使用即时编译（JIT）进行编译。

JIT 编译可以提高某些模型的性能。然而，它可能不适用于所有 TensorFlow 操作。如果您使用的是自定义模型或层，并且看到与 XLA 相关的错误，您可能需要将 jit_compile 参数设置为 False。以下是使用 XLA 和 TensorFlow 时遇到的已知问题的列表。除了这些问题，还有一些操作不受 XLA 支持。

您可能遇到的错误信息如下：

Detected unsupported operations when trying to compile graph
__inference_one_step_on_data_125[] on XLA_GPU_JIT

例如，以下代码片段将重现上述错误：

class MyModel(keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def call(self, inputs):
        string_input = tf.strings.as_string(inputs)
        return tf.strings.to_number(string_input)


subclass_model = MyModel()
x_train = np.array([[1, 2, 3], [4, 5, 6]])
subclass_model.compile(optimizer="sgd", loss="mse")
subclass_model.predict(x_train)

如何修复: 在 model.compile(..., jit_compile=False) 中设置 jit_compile=False，或者将 jit_compile 属性设置为 False，如下所示：

class MyModel(keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def call(self, inputs):
        # tf.strings 操作不被 XLA 支持
        string_input = tf.strings.as_string(inputs)
        return tf.strings.to_number(string_input)


subclass_model = MyModel()
x_train = np.array([[1, 2, 3], [4, 5, 6]])
subclass_model.jit_compile = False
subclass_model.predict(x_train)

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 51ms/step

array([[1., 2., 3.],
       [4., 5., 6.]], dtype=float32)

以 TF SavedModel 格式保存模型

通过 model.save() 保存为 TF SavedModel 格式在 Keras 3 中不再受支持。

您可能会遇到的错误信息如下：

>>> model.save("mymodel")
ValueError: 保存的文件路径扩展无效。请添加 `.keras` 扩展名以使用原生 Keras 格式（推荐）或 `.h5` 扩展名。如果您想导出用于 TFLite/TFServing/等的 SavedModel，请使用 `model.export(filepath)`。收到: filepath=saved_model.

以下代码片段将重现上述错误：

sequential_model = keras.Sequential([
    keras.layers.Dense(2)
])
sequential_model.save("saved_model")

如何修复： 使用 model.export(filepath) 代替 model.save(filepath)

sequential_model = keras.Sequential([keras.layers.Dense(2)])
sequential_model(np.random.rand(3, 5))
sequential_model.export("saved_model")

INFO:tensorflow:资产写入: saved_model/assets

INFO:tensorflow:资产写入: saved_model/assets

已保存工件于 'saved_model'。以下端点可用：

* 端点 'serve'
  args_0 (POSITIONAL_ONLY): TensorSpec(shape=(3, 5), dtype=tf.float32, name='keras_tensor')
输出类型:
  TensorSpec(shape=(3, 2), dtype=tf.float32, name=None)
捕获:
  14428321600: TensorSpec(shape=(), dtype=tf.resource, name=None)
  14439128528: TensorSpec(shape=(), dtype=tf.resource, name=None)

加载 TF SavedModel

通过 keras.models.load_model() 加载 TF SavedModel 文件不再受支持如果您尝试使用 keras.models.load_model() 加载 TF SavedModel，将会出现以下错误：

ValueError: 文件格式不受支持: filepath=saved_model。 Keras 3 只支持 V3
`.keras` 文件和遗留的 H5 格式文件（`.h5` 扩展名）。 请注意，遗留
SavedModel 格式在 Keras 3 中不被 `load_model()` 支持。 为了重新加载
TensorFlow SavedModel 作为 Keras 3 中的推理层，请使用
`keras.layers.TFSMLayer(saved_model, call_endpoint='serving_default')`（请注意，您的
`call_endpoint` 可能具有不同的名称）。

以下代码片段将重现上述错误：

keras.models.load_model("saved_model")

如何修复： 使用 keras.layers.TFSMLayer(filepath, call_endpoint="serving_default") 重新加载 TF SavedModel 作为 Keras 层。这不限于源自 Keras 的 SavedModels – 它将适用于任何 SavedModel，例如 TF-Hub 模型。

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

<TFSMLayer name=tfsm_layer, built=True>

在功能模型中使用深度嵌套输入

Model() 不再可以传递深度嵌套的输入/输出（嵌套超过 1 层深，例如张量的列表的列表）。

您将遇到如下错误：

ValueError: 当提供 `inputs` 作为字典时，字典中的所有值必须为
KerasTensors。 收到: inputs={'foo': <KerasTensor shape=(None, 1), dtype=float32,
sparse=None, name=foo>, 'bar': {'baz': <KerasTensor shape=(None, 1), dtype=float32,
sparse=None, name=bar>}} 包含无效值 {'baz': <KerasTensor shape=(None, 1),
dtype=float32, sparse=None, name=bar>} 的类型 <class 'dict'>

以下代码片段将重现上述错误：

inputs = {
    "foo": keras.Input(shape=(1,), name="foo"),
    "bar": {
        "baz": keras.Input(shape=(1,), name="bar"),
    },
}
outputs = inputs["foo"] + inputs["bar"]["baz"]
keras.Model(inputs, outputs)

如何修复： 用输入张量的字典、列表和元组替换嵌套输入。

inputs = {
    "foo": keras.Input(shape=(1,), name="foo"),
    "bar": keras.Input(shape=(1,), name="bar"),
}
outputs = inputs["foo"] + inputs["bar"]
keras.Model(inputs, outputs)

<Functional name=functional_2, built=True>

TF 自动图

在 Keras 2 中，TF 自动图默认在自定义层的 call() 方法上启用。在 Keras 3 中，它不再启用。这意味着如果您使用控制流，您可能需要使用条件操作，或者您可以将 call() 方法装饰为 @tf.function。

您将遇到如下错误：

OperatorNotAllowedInGraphError: 调用 MyCustomLayer.call() 时遇到异常。

使用符号 [`tf.Tensor`](https://www.tensorflow.org/api_docs/python/tf/Tensor) 作为 Python `bool` 是不允许的。 您可以尝试以下解决方案来解决此问题：如果您在图模式下运行，请使用急切执行模式或将此函数装饰为 @tf.function。如果您使用 AutoGraph，您可以尝试将此函数装饰为 @tf.function。如果这样仍然不行，那么您可能在使用不受支持的特性，或者您的源代码可能对 AutoGraph 不可见。
这里有一个 [链接以获取更多信息](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md#access-to-source-code)。

以下代码片段将重现上述错误：

class MyCustomLayer(keras.layers.Layer):

  def call(self, inputs):
    if tf.random.uniform(()) > 0.5:
      return inputs * 2
    else:
      return inputs / 2


layer = MyCustomLayer()
data = np.random.uniform(size=[3, 3])
model = keras.models.Sequential([layer])
model.compile(optimizer="adam", loss="mse")
model.predict(data)

如何修复： 使用 @tf.function 装饰器装饰你的 call() 方法

class MyCustomLayer(keras.layers.Layer):
    @tf.function()
    def call(self, inputs):
        if tf.random.uniform(()) > 0.5:
            return inputs * 2
        else:
            return inputs / 2


layer = MyCustomLayer()
data = np.random.uniform(size=[3, 3])
model = keras.models.Sequential([layer])
model.compile(optimizer="adam", loss="mse")
model.predict(data)

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 43ms/step

array([[0.59727275, 1.9986179 , 1.5514829 ],
       [0.56239295, 1.6529864 , 0.33085832],
       [0.67086476, 1.5208522 , 1.99276   ]], dtype=float32)

使用 `KerasTensor` 调用 TF 操作

在功能模型构建期间，禁止将 TF 操作应用于 Keras 张量：“KerasTensor 不能作为 TensorFlow 函数的输入”。

你会遇到的错误如下：

ValueError: A KerasTensor cannot be used as input to a TensorFlow function. A KerasTensor
is a symbolic placeholder for a shape and dtype, used when constructing Keras Functional
models or Keras Functions. You can only use it as input to a Keras layer or a Keras
operation (from the namespaces `keras.layers` and `keras.operations`).

以下代码片段将重现该错误：

input = keras.layers.Input([2, 2, 1])
tf.squeeze(input)

如何修复： 使用 keras.ops 中的等效操作。

input = keras.layers.Input([2, 2, 1])
keras.ops.squeeze(input)

<KerasTensor shape=(None, 2, 2), dtype=float32, sparse=None, name=keras_tensor_6>

多输出模型的 `evaluate()`

多输出模型的 evaluate() 方法不再单独返回每个输出损失。相反，你应该在 compile() 方法中使用 metrics 参数来跟踪这些损失。

在处理多个命名输出（如 output_a 和 output_b）时，传统的 tf.keras 将在 metrics 中包含 _loss、_loss 和类似条目。然而，在 keras 3.0 中，这些条目不会自动添加到 metrics 中。必须在每个单独输出的 metrics 列表中显式提供它们。

以下代码片段会重现上述行为：

from keras import layers
# 一个具有多个输出的功能模型
inputs = layers.Input(shape=(10,))
x1 = layers.Dense(5, activation='relu')(inputs)
x2 = layers.Dense(5, activation='relu')(x1)
output_1 = layers.Dense(5, activation='softmax', name="output_1")(x1)
output_2 = layers.Dense(5, activation='softmax', name="output_2")(x2)
model = keras.Model(inputs=inputs, outputs=[output_1, output_2])
model.compile(optimizer='adam', loss='categorical_crossentropy')
# 虚拟数据
x_test = np.random.uniform(size=[10, 10])
y_test = np.random.uniform(size=[10, 5])

model.evaluate(x_test, y_test)

from keras import layers

# 一个具有多个输出的功能模型
inputs = layers.Input(shape=(10,))
x1 = layers.Dense(5, activation="relu")(inputs)
x2 = layers.Dense(5, activation="relu")(x1)
output_1 = layers.Dense(5, activation="softmax", name="output_1")(x1)
output_2 = layers.Dense(5, activation="softmax", name="output_2")(x2)
# 虚拟数据
x_test = np.random.uniform(size=[10, 10])
y_test = np.random.uniform(size=[10, 5])
multi_output_model = keras.Model(inputs=inputs, outputs=[output_1, output_2])
multi_output_model.compile(
    optimizer="adam",
    loss="categorical_crossentropy",
    metrics=["categorical_crossentropy", "categorical_crossentropy"],
)
multi_output_model.evaluate(x_test, y_test)

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 112ms/step - loss: 4.0217 - output_1_categorical_crossentropy: 4.0217

[4.021683692932129, 4.021683692932129]

TensorFlow 变量跟踪

将 tf.Variable 设置为 Keras 3 层或模型的属性将不会自动跟踪该变量，这与 Keras 2 不同。以下代码片段将显示 tf.Variable 没有被跟踪。

class MyCustomLayer(keras.layers.Layer):
    def __init__(self, units):
        super().__init__()
        self.units = units

    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.w = tf.Variable(initial_value=tf.zeros([input_dim, self.units]))
        self.b = tf.Variable(initial_value=tf.zeros([self.units,]))

    def call(self, inputs):
        return keras.ops.matmul(inputs, self.w) + self.b


layer = MyCustomLayer(3)
data = np.random.uniform(size=[3, 3])
model = keras.models.Sequential([layer])
model.compile(optimizer="adam", loss="mse")
model.predict(data)
# 模型没有任何可训练的变量
for layer in model.layers:
    print(layer.trainable_variables)

你将看到如下警告：

UserWarning: 该模型没有任何可训练的权重。
  warnings.warn("该模型没有任何可训练的权重。")

如何解决： 使用 self.add_weight() 方法或选择 keras.Variable。如果您当前正在使用 tf.variable，可以切换到 keras.Variable。

class MyCustomLayer(keras.layers.Layer):
    def __init__(self, units):
        super().__init__()
        self.units = units

    def build(self, input_shape):
        input_dim = input_shape[-1]
        self.w = self.add_weight(
            shape=[input_dim, self.units],
            initializer="zeros",
        )
        self.b = self.add_weight(
            shape=[
                self.units,
            ],
            initializer="zeros",
        )

    def call(self, inputs):
        return keras.ops.matmul(inputs, self.w) + self.b


layer = MyCustomLayer(3)
data = np.random.uniform(size=[3, 3])
model = keras.models.Sequential([layer])
model.compile(optimizer="adam", loss="mse")
model.predict(data)
# 验证变量现在是否正在被跟踪
for layer in model.layers:
    print(layer.trainable_variables)

 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 33ms/step
[<KerasVariable shape=(3, 3), dtype=float32, path=sequential_2/my_custom_layer_1/variable>, <KerasVariable shape=(3,), dtype=float32, path=sequential_2/my_custom_layer_1/variable_1>]

嵌套 `call()` 参数中的 `None` 条目

在 Layer.call() 的嵌套（例如列表/元组）张量参数中不允许出现 None 条目，也不允许在 call() 的嵌套返回值中出现 None。

如果参数中的 None 是故意的并且有特定的目的，确保该参数是可选的并将其结构化为单独的参数。例如，考虑将 call 方法定义为可选参数。

以下代码片段将重现该错误。

class CustomLayer(keras.layers.Layer):
    def __init__(self):
        super().__init__()

    def call(self, inputs):
        foo = inputs["foo"]
        baz = inputs["bar"]["baz"]
        if baz is not None:
            return foo + baz
        return foo

layer = CustomLayer()
inputs = {
    "foo": keras.Input(shape=(1,), name="foo"),
    "bar": {
        "baz": None,
    },
}
layer(inputs)

如何解决：

解决方案 1: 用一个值替换 None，如下所示：

class CustomLayer(keras.layers.Layer):
    def __init__(self):
        super().__init__()

    def call(self, inputs):
        foo = inputs["foo"]
        baz = inputs["bar"]["baz"]
        return foo + baz


layer = CustomLayer()
inputs = {
    "foo": keras.Input(shape=(1,), name="foo"),
    "bar": {
        "baz": keras.Input(shape=(1,), name="bar"),
    },
}
layer(inputs)

<KerasTensor shape=(None, 1), dtype=float32, sparse=False, name=keras_tensor_14>

解决方案 2: 使用可选参数定义 call 方法。下面是此修复的示例：

class CustomLayer(keras.layers.Layer):
    def __init__(self):
        super().__init__()

    def call(self, foo, baz=None):
        if baz is not None:
            return foo + baz
        return foo


layer = CustomLayer()
foo = keras.Input(shape=(1,), name="foo")
baz = None
layer(foo, baz=baz)

<KerasTensor shape=(None, 1), dtype=float32, sparse=False, name=keras_tensor_15>

状态构建问题

Keras 3 比 Keras 2 在何时可以创建状态（例如数值权重变量）方面严格得多。Keras 3 希望所有状态在模型可以训练之前创建。这是使用 JAX 的要求（而 TensorFlow 对状态创建时机非常宽松）。

Keras 层应该在它们的构造函数（__init__() 方法）或它们的 build() 方法中创建其状态。它们应该避免在 call() 中创建状态。

如果您忽略此建议并在 call() 中创建状态（例如，通过调用一个之前未构建的层），则 Keras 将尝试通过在训练之前对符号输入调用 call() 方法来自动构建该层。然而，这种自动状态创建的尝试在某些情况下可能会失败。这会导致错误，内容如下：

Layer 'frame_position_embedding' looks like it has unbuilt state,
but Keras is not able to trace the layer `call()` in order to build it automatically.
Possible causes:
1. The `call()` method of your layer may be crashing.
Try to `__call__()` the layer eagerly on some test input first to see if it works.
E.g. `x = np.random.random((3, 4)); y = layer(x)`
2. If the `call()` method is correct, then you may need to implement
the `def build(self, input_shape)` method on your layer.
It should create all variables used by the layer
(e.g. by calling `layer.build()` on all its children layers).

您可以通过以下层重现此错误，当与 JAX 后端一起使用时：

class PositionalEmbedding(keras.layers.Layer):
    def __init__(self, sequence_length, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.position_embeddings = layers.Embedding(
            input_dim=sequence_length, output_dim=output_dim
        )
        self.sequence_length = sequence_length
        self.output_dim = output_dim

    def call(self, inputs):
        inputs = keras.ops.cast(inputs, self.compute_dtype)
        length = keras.ops.shape(inputs)[1]
        positions = keras.ops.arange(start=0, stop=length, step=1)
        embedded_positions = self.position_embeddings(positions)
        return inputs + embedded_positions  # 返回输入加上嵌入的位置

如何修复它： 完全按照错误信息的要求进行操作。首先，尝试以急切方式运行该层，以检查 call() 方法是否确实正确（注意：如果它在 Keras 2 中有效，那么它就是正确的，并且不需要更改）。如果确实正确，那么你应该实现一个 build(self, input_shape) 方法，用于创建该层的所有状态，包括子层的状态。以下是为上面的层应用的修复（注意 build() 方法）：

class PositionalEmbedding(keras.layers.Layer):
    def __init__(self, sequence_length, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.position_embeddings = layers.Embedding(
            input_dim=sequence_length, output_dim=output_dim
        )
        self.sequence_length = sequence_length
        self.output_dim = output_dim

    def build(self, input_shape):
        self.position_embeddings.build(input_shape)

    def call(self, inputs):
        inputs = keras.ops.cast(inputs, self.compute_dtype)
        length = keras.ops.shape(inputs)[1]
        positions = keras.ops.arange(start=0, stop=length, step=1)
        embedded_positions = self.position_embeddings(positions)
        return inputs + embedded_positions

移除的特性

在 Keras 3 中，少量使用率非常低的遗留特性被移除，以进行清理：

keras.layers.ThresholdedReLU 被移除。相反，你可以简单地使用带有 threshold 参数的 ReLU 层。
符号性的 Layer.add_loss()：符号性的 add_loss() 被移除（你仍然可以在层/模型的 call() 方法中使用 add_loss()）。
由于使用率非常低，局部连接层（LocallyConnected1D，LocallyConnected2D）被移除。要使用局部连接层，请将层的实现复制到你自己的代码库中。
由于使用率非常低，keras.layers.experimental.RandomFourierFeatures 被移除。要使用它，请将层的实现复制到你自己的代码库中。
移除的层属性：层属性 metrics，dynamic 被移除。metrics 在 Model 类中仍然可用。
RNN 层中的 constants 和 time_major 参数被移除。constants 参数是 Theano 的遗留物且使用率非常低。time_major 参数的使用率也很低。
reset_metrics 参数：reset_metrics 参数从 model.*_on_batch() 方法中移除。此参数的使用率非常低。
keras.constraints.RadialConstraint 对象被移除。此对象的使用率非常低。

过渡到与后端无关的 Keras 3

使用 TensorFlow 后端的 Keras 3 代码将与原生的 TensorFlow API 一起工作。然而，如果你希望你的代码与后端无关，你需要：

用等效的 Keras API 替换所有的 tf.* API 调用。
将你的自定义 train_step/test_step 方法转换为多框架实现。
确保你在层中正确使用无状态的 keras.random ops。

我们来详细讨论每一点。

切换到 Keras ops

在许多情况下，这可能是你开始能够使用 JAX 和 PyTorch 运行自定义层和指标所需做的唯一事情：将所有的 tf.*，tf.math*，tf.linalg.* 等替换为 keras.ops.*。大多数 TF ops 应与 Keras 3 保持一致。如果名称不同，它们将在本指南中予以强调。

NumPy ops

Keras 将 NumPy API 实现为 keras.ops 的一部分。

下表仅列出了少量 TensorFlow 和 Keras ops；未列出的 ops 在两个框架中通常名称相同（例如 reshape，matmul，cast 等）。

TensorFlow	Keras 3.0
`tf.abs`	`keras.ops.absolute`
`tf.reduce_all`	`keras.ops.all`
`tf.reduce_max`	`keras.ops.amax`
`tf.reduce_min`	`keras.ops.amin`
`tf.reduce_any`	`keras.ops.any`
`tf.concat`	`keras.ops.concatenate`
`tf.range`	`keras.ops.arange`
`tf.acos`	`keras.ops.arccos`
`tf.asin`	`keras.ops.arcsin`
`tf.asinh`	`keras.ops.arcsinh`
`tf.atan`	`keras.ops.arctan`
`tf.atan2`	`keras.ops.arctan2`
`tf.atanh`	`keras.ops.arctanh`
`tf.convert_to_tensor`	`keras.ops.convert_to_tensor`
`tf.reduce_mean`	`keras.ops.mean`
`tf.clip_by_value`	`keras.ops.clip`
`tf.math.conj`	`keras.ops.conjugate`
`tf.linalg.diag_part`	`keras.ops.diagonal`
`tf.reverse`	`keras.ops.flip`
`tf.gather`	`keras.ops.take`
`tf.math.is_finite`	`keras.ops.isfinite`
`tf.math.is_inf`	`keras.ops.isinf`
`tf.math.is_nan`	`keras.ops.isnan`
`tf.reduce_max`	`keras.ops.max`
`tf.reduce_mean`	`keras.ops.mean`
`tf.reduce_min`	`keras.ops.min`
`tf.rank`	`keras.ops.ndim`
`tf.math.pow`	`keras.ops.power`
`tf.reduce_prod`	`keras.ops.prod`
`tf.math.reduce_std`	`keras.ops.std`
`tf.reduce_sum`	`keras.ops.sum`
`tf.gather`	`keras.ops.take`
`tf.gather_nd`	`keras.ops.take_along_axis`
`tf.math.reduce_variance`	`keras.ops.var`

其他操作

TensorFlow	Keras 3.0
`tf.nn.sigmoid_cross_entropy_with_logits`	`keras.ops.binary_crossentropy` (注意 `from_logits` 参数)
`tf.nn.sparse_softmax_cross_entropy_with_logits`	`keras.ops.sparse_categorical_crossentropy` (注意 `from_logits` 参数)
`tf.nn.sparse_softmax_cross_entropy_with_logits`	`keras.ops.categorical_crossentropy(target, output, from_logits=False, axis=-1)`
`tf.nn.conv1d`, `tf.nn.conv2d`, `tf.nn.conv3d`, `tf.nn.convolution`	`keras.ops.conv`
`tf.nn.conv_transpose`, `tf.nn.conv1d_transpose`, `tf.nn.conv2d_transpose`, `tf.nn.conv3d_transpose`	`keras.ops.conv_transpose`
`tf.nn.depthwise_conv2d`	`keras.ops.depthwise_conv`
`tf.nn.separable_conv2d`	`keras.ops.separable_conv`
`tf.nn.batch_normalization`	没有直接等效; 使用 `keras.layers.BatchNormalization`
`tf.nn.dropout`	`keras.random.dropout`
`tf.nn.embedding_lookup`	`keras.ops.take`
`tf.nn.l2_normalize`	`keras.utils.normalize` (不是操作)
`x.numpy`	`keras.ops.convert_to_numpy`
`tf.scatter_nd_update`	`keras.ops.scatter_update`
`tf.tensor_scatter_nd_update`	`keras.ops.slice_update`
`tf.signal.fft2d`	`keras.ops.fft2`
`tf.signal.inverse_stft`	`keras.ops.istft`

自定义 `train_step()` 方法

您的模型可能包括自定义的 train_step() 或 test_step() 方法，这些方法依赖于仅限 TensorFlow 的 API – 例如，您的 train_step() 方法可能利用 TensorFlow 的 tf.GradientTape。要将这些模型转换为在 JAX 或 PyTorch 上运行，您需要为每个您想要支持的后端编写不同的 train_step() 实现。

在某些情况下，您可能能够简单地重写 Model.compute_loss() 方法，使其完全与后端无关，而不是重写 train_step()。以下是一个具有自定义 compute_loss() 方法的层的示例，该方法在 JAX、TensorFlow 和 PyTorch 中都可以工作：

class MyModel(keras.Model):
    def compute_loss(self, x=None, y=None, y_pred=None, sample_weight=None):
        loss = keras.ops.sum(keras.losses.mean_squared_error(y, y_pred, sample_weight))
        return loss

如果您需要修改优化机制本身，而不仅仅是损失计算，那么您需要重写 train_step()，并为每个后端实现一个 train_step 方法，如下所示。

有关如何处理每个后端的详细信息，请参阅以下指南：

使用 JAX 自定义 fit() 中发生的操作
使用 TensorFlow 自定义 fit() 中发生的操作
使用 PyTorch 自定义 fit() 中发生的操作

class MyModel(keras.Model):
    def train_step(self, *args, **kwargs):
        if keras.backend.backend() == "jax":
            return self._jax_train_step(*args, **kwargs)
        elif keras.backend.backend() == "tensorflow":
            return self._tensorflow_train_step(*args, **kwargs)
        elif keras.backend.backend() == "torch":
            return self._torch_train_step(*args, **kwargs)

    def _jax_train_step(self, state, data):
        pass  # 查看指南: keras.io/guides/custom_train_step_in_jax/

    def _tensorflow_train_step(self, data):
        pass  # 查看指南: keras.io/guides/custom_train_step_in_tensorflow/

    def _torch_train_step(self, data):
        pass  # 查看指南: keras.io/guides/custom_train_step_in_torch/

使用 RNG 的层

Keras 3 具有新的 keras.random 命名空间，包含：

这些操作是 无状态的，这意味着如果您传递 seed 参数，它们每次将返回相同的结果。像这样：

print(keras.random.normal(shape=(), seed=123))
print(keras.random.normal(shape=(), seed=123))

tf.Tensor(0.7832616, shape=(), dtype=float32)
tf.Tensor(0.7832616, shape=(), dtype=float32)

关键是，这与有状态的 tf.random 操作的行为不同：

print(tf.random.normal(shape=(), seed=123))
print(tf.random.normal(shape=(), seed=123))

tf.Tensor(2.4435377, shape=(), dtype=float32)
tf.Tensor(-0.6386405, shape=(), dtype=float32)

当您编写使用 RNG 的层，例如自定义 dropout 层时，您会希望在层调用时使用不同的种子值。然而，您不能只增加一个 Python 整数并将其传递，因为虽然这在执行时会正常工作，但在使用编译时将无法按预期工作（编译可用于 JAX、TensorFlow 和 PyTorch）。当编译层时，层首次遇到的 Python 整数种子值将被硬编码到编译图中。

为了解决这个问题，您应该将一个状态ful 的 keras.random.SeedGenerator 对象的实例作为 seed 参数传递，如下所示：

seed_generator = keras.random.SeedGenerator(1337)
print(keras.random.normal(shape=(), seed=seed_generator))
print(keras.random.normal(shape=(), seed=seed_generator))

tf.Tensor(0.6077996, shape=(), dtype=float32)
tf.Tensor(0.8211102, shape=(), dtype=float32)

因此，在编写一个使用 RNG 的层时，您将使用以下模式：

class RandomNoiseLayer(keras.layers.Layer):
    def __init__(self, noise_rate, **kwargs):
        super().__init__(**kwargs)
        self.noise_rate = noise_rate
        self.seed_generator = keras.random.SeedGenerator(1337)

    def call(self, inputs):
        noise = keras.random.uniform(
            minval=0, maxval=self.noise_rate, seed=self.seed_generator
        )
        return inputs + noise

这样的层在任何环境中都可以安全使用——在急切执行或编译模型中。每次层调用将使用不同的种子值，正如预期的那样。

将 Keras 2 代码迁移到多后端 Keras 3

◆ 设置

◆ 从 Keras 2 转到 Keras 3（使用 TensorFlow 后端）

在 GPU 上，jit_compile 默认设置为 True。

以 TF SavedModel 格式保存模型

加载 TF SavedModel

在功能模型中使用深度嵌套输入

TF 自动图

使用 KerasTensor 调用 TF 操作

多输出模型的 evaluate()

TensorFlow 变量跟踪

嵌套 call() 参数中的 None 条目

状态构建问题

移除的特性

◆ 过渡到与后端无关的 Keras 3

切换到 Keras ops

自定义 train_step() 方法

使用 RNG 的层

将 Keras 2 代码迁移到多后端 Keras 3

设置

从 Keras 2 转到 Keras 3（使用 TensorFlow 后端）

在 GPU 上，jit_compile 默认设置为 True。

以 TF SavedModel 格式保存模型

加载 TF SavedModel

在功能模型中使用深度嵌套输入

TF 自动图

使用 KerasTensor 调用 TF 操作

多输出模型的 evaluate()

TensorFlow 变量跟踪

嵌套 call() 参数中的 None 条目

状态构建问题

移除的特性

过渡到与后端无关的 Keras 3

切换到 Keras ops

NumPy ops

其他操作

自定义 train_step() 方法

使用 RNG 的层

在 GPU 上，`jit_compile` 默认设置为 `True`。

使用 `KerasTensor` 调用 TF 操作

多输出模型的 `evaluate()`

嵌套 `call()` 参数中的 `None` 条目

自定义 `train_step()` 方法