加载Keras Sequential模型时遇TypeError：意外参数'reduction'

问题 #

加载Keras Sequential模型时，触发TypeError: __init__() got an unexpected keyword argument 'reduction'错误。相关代码及错误信息如下：

训练的模型结构 #

model = Sequential([
    tf.keras.Input(shape=(in_dim,)),
    layers.Dense(
        units=(in_dim+1),
        activation=layers.LeakyReLU(alpha=.01)
    ),
    layers.Dropout(rate=.05),
    layers.Dense(
        units=(in_dim),
        activation=layers.LeakyReLU(alpha=.01)
    ),
    layers.Dropout(rate=.05),
    layers.Dense(units=1, activation="sigmoid")
])

编译与拟合代码 #

def compile_and_fit(
    model,
    name,
    X_y_train,
    X_y_val,
    optimizer=None,
    max_epochs=10000,
    batch_size=BATCH_SIZE,
):

    X_train, y_train = X_y_train
    steps_per_epoch = len(X_train) // batch_size
    X_val, y_val = X_y_val
    steps_per_epoch_val = len(X_val) // batch_size

    if optimizer is None:
        optimizer = get_optimizer(steps_per_epoch)
    model.compile(
        optimizer=optimizer,
        loss=tf.keras.losses.BinaryCrossentropy(),
        metrics=[
            tf.keras.losses.BinaryCrossentropy(name="binary_crossentropy"),
            tf.keras.metrics.Precision(name="precision"),
            "accuracy",
        ],
    )

    model.summary()

    history = model.fit(
        x=X_train,
        y=y_train,
        steps_per_epoch=steps_per_epoch,
        batch_size=batch_size,
        epochs=max_epochs,
        validation_data=X_y_val,
        validation_steps=steps_per_epoch_val,
        callbacks=get_callbacks(name),
        verbose=1,
    )
    return history

模型保存与加载代码 #

保存：

model.save(f"./saved_model/my_model", save_format="tf")

加载：

model = tf.keras.models.load_model("./saved_model/my_model")

错误栈 #

TypeError                                 Traceback (most recent call last)
Cell In [25], line 2
----> 2 model = tf.keras.models.load_model("./saved_model/my_model")

File /mnt/c/Code/venv/lib/python3.9/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File /mnt/c/Code/venv/lib/python3.9/site-packages/keras/dtensor/utils.py:144, in inject_mesh.<locals>._wrap_function(instance, *args, **kwargs)
    142 if mesh is not None:
    143     instance._mesh = mesh
--> 144 init_method(instance, *args, **kwargs)

TypeError: __init__() got an unexpected keyword argument 'reduction'

修复思路 #

核心原因 #

问题出在将损失函数类（tf.keras.losses.BinaryCrossentropy）直接作为指标传入metrics列表。损失函数类的初始化逻辑和指标类不同：加载模型时，Keras会自动给指标对象传入reduction参数，但tf.keras.losses.BinaryCrossentropy的构造函数并不接受该参数，因此触发错误。

具体修复步骤 #

1. 替换指标为对应指标类
   将编译代码中metrics列表里的tf.keras.losses.BinaryCrossentropy(name="binary_crossentropy")替换为tf.keras.metrics.BinaryCrossentropy(name="binary_crossentropy")。修改后的编译部分代码：

   model.compile(
       optimizer=optimizer,
       loss=tf.keras.losses.BinaryCrossentropy(),
       metrics=[
           tf.keras.metrics.BinaryCrossentropy(name="binary_crossentropy"),
           tf.keras.metrics.Precision(name="precision"),
           "accuracy",
       ],
   )
   

2. 重新训练并保存模型
   用修改后的代码重新训练模型，再执行保存操作。此时保存的模型加载时不会再触发reduction参数错误。
3. （可选）无需重新训练的临时修复
   如果无法重新训练，加载模型时可以通过custom_objects自定义损失函数的初始化逻辑，忽略reduction参数：

   from keras.losses import BinaryCrossentropy
   
   def custom_binary_crossentropy(**kwargs):
       # 移除reduction参数后初始化损失函数
       kwargs.pop('reduction', None)
       return BinaryCrossentropy(**kwargs)
   
   model = tf.keras.models.load_model(
       "./saved_model/my_model",
       custom_objects={'BinaryCrossentropy': custom_binary_crossentropy}
   )
   

   注意：这种方法属于临时兼容方案，更推荐使用指标类的规范写法。

内容的提问来源于stack exchange，提问作者Gil Ben David