tf.keras自定义含额外参数损失函数报错及跨框架适配方案

为什么会报错？ #

在原生Keras 2.2.4里，你的代码能正常运行是因为框架不会主动把sample_weight参数传递给自定义损失函数——毕竟你是把权重作为模型的第二个输入张量来处理的。但换到tf.keras 2.2.4-tf时，框架的实现逻辑变了：它会强制把sample_weight作为关键字参数传给损失函数，不管你有没有用到这个参数。而你的example_loss函数根本没定义这个参数，自然就触发了TypeError: example_loss() got an unexpected keyword argument 'sample_weight'。

快速修复：适配tf.keras的修改 #

最简单的办法就是给你的自定义损失函数加一个可选的sample_weight参数，哪怕你不用它也行：

def example_loss(y_true, y_pred, sample_weights_, sample_weight=None):
    # 忽略框架自动传入的sample_weight，用我们预先定义的sample_weights_
    per_sample_loss = K.sqrt(K.sum(K.pow(y_pred - y_true, 2), axis=-1))
    return K.mean(per_sample_loss * K.flatten(sample_weights_))

修改后，tf.keras传递sample_weight参数时，函数有对应的可选参数接收，就不会报错了，原生Keras也能正常运行。

更优雅的兼容方案：使用框架标准的sample_weight参数 #

其实不管是原生Keras还是tf.keras，.fit方法的sample_weight参数就是专门用来处理样本权重的，完全没必要把权重作为模型的输入。这种写法不仅更符合框架设计，还能完美兼容两个环境：

import numpy as np
from tensorflow.keras import Input, Dense, Model
from tensorflow.keras.optimizers import Adam
import tensorflow.keras.backend as K

# 生成模拟数据
train_X = np.random.randn(100, 5)
train_Y = np.random.randn(100, 1)
# sample_weight需要是一维数组，所以这里把它展平
train_sample_weights = np.random.randn(*train_Y.shape).flatten()

# 自定义损失函数：返回每个样本的损失（框架会自动用sample_weight加权）
def example_loss(y_true, y_pred):
    # 计算单个样本的损失，shape为(batch_size,)
    return K.sqrt(K.sum(K.pow(y_pred - y_true, 2), axis=-1))

# 构建模型（不需要权重输入）
input_tensor = Input(shape=(train_X.shape[1],))
x = Dense(100, activation="relu")(input_tensor)
out = Dense(1)(x)
model = Model(input_tensor, out)

# 编译模型
model.compile(optimizer="Adam", loss=example_loss)

# 训练时直接传入sample_weight参数
model.fit(x=train_X, y=train_Y, sample_weight=train_sample_weights, epochs=10)

框架会自动将每个样本的损失乘以对应的权重，再计算加权平均，完全不需要你手动处理权重的传递和计算。

如果你坚持要把权重作为模型输入... #

如果你的业务逻辑必须把权重作为模型的输入张量（比如权重需要参与网络的前向计算），那除了给损失函数加可选参数外，还要注意不要在.fit里传入sample_weight参数，避免重复传递：

from functools import partial
import numpy as np
from tensorflow.keras import Input, Dense, Model
from tensorflow.keras.optimizers import Adam
import tensorflow.keras.backend as K

# 生成模拟数据
train_X = np.random.randn(100, 5)
train_Y = np.random.randn(100, 1)
train_sample_weights = np.random.randn(*train_Y.shape)

# 带可选sample_weight参数的损失函数
def example_loss(y_true, y_pred, sample_weights_, sample_weight=None):
    per_sample_loss = K.sqrt(K.sum(K.pow(y_pred - y_true, 2), axis=-1))
    return K.mean(per_sample_loss * K.flatten(sample_weights_))

# 构建带权重输入的模型
input_tensor = Input(shape=(train_X.shape[1],))
weights_tensor = Input(shape=(train_sample_weights.shape[1],))
x = Dense(100, activation="relu")(input_tensor)
out = Dense(1)(x)

# 绑定sample_weights_参数
loss_function = partial(example_loss, sample_weights_=weights_tensor)
model = Model([input_tensor, weights_tensor], out)

# 编译模型
model.compile(optimizer="Adam", loss=loss_function)

# 训练时只传入数据和自定义权重，不要传sample_weight参数
model.fit(x=[train_X, train_sample_weights], y=train_Y, epochs=10)

这样修改后，两个框架环境下都能正常运行。

内容的提问来源于stack exchange，提问作者F1sher