TensorFlow含自定义损失函数的LSTM模型保存加载报错求助

问题根源 #

你遇到的ValueError: Unknown loss function: 'nested_loss'错误，是因为调用composite_loss(alpha)时返回的是内部嵌套的nested_loss函数，模型保存时会将损失函数标识为nested_loss，但你仅注册了外层的composite_loss，导致加载时无法找到对应的损失函数。

可靠解决方案：使用Keras Loss类实现 #

用类继承tf.keras.losses.Loss来定义损失函数，是Keras官方推荐的自定义损失序列化方式，能避免嵌套函数带来的序列化问题。

1. 重写自定义损失类 #

import tensorflow as tf
from tensorflow.keras.losses import Loss

# 自定义绝对值求和损失
class CustomLoss(Loss):
    def call(self, y_true, y_pred):
        return tf.reduce_sum(tf.abs(y_true - y_pred))

# 复合损失类：结合MSE和自定义损失
class CompositeLoss(Loss):
    def __init__(self, alpha=0.5, name="composite_loss"):
        super().__init__(name=name)
        self.alpha = alpha
        self.mse_loss = tf.keras.losses.MeanSquaredError()
        self.custom_loss = CustomLoss()

    def call(self, y_true, y_pred):
        mse_val = self.mse_loss(y_true, y_pred)
        custom_val = self.custom_loss(y_true, y_pred)
        return self.alpha * mse_val + custom_val

2. 模型编译、保存 #

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dropout, Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.keras import utils

# 构建LSTM模型
model = Sequential()
model.add(LSTM(128, return_sequences=True, input_shape=(X_train.shape[1], 1)))
model.add(Dropout(0.3))
model.add(LSTM(64, return_sequences=True))
model.add(Dropout(0.3))
model.add(LSTM(32, return_sequences=True))
model.add(Dropout(0.3))
model.add(LSTM(16, return_sequences=False))
model.add(Dense(1))

# 编译模型：传入复合损失实例
alpha = 0.5
optimizer = Adam(learning_rate=0.001)  # 新版本TensorFlow用learning_rate替代lr
model.compile(loss=CompositeLoss(alpha), optimizer=optimizer)

# 注册自定义损失类，确保保存时能被序列化
utils.get_custom_objects()['CompositeLoss'] = CompositeLoss
utils.get_custom_objects()['CustomLoss'] = CustomLoss

# 保存模型
model.save('SYM-Revised-CAS-30')

3. 模型加载 #

from tensorflow.keras.models import load_model

# 加载时注册对应的自定义损失类
model = load_model('SYM-Revised-CAS-30', 
                   custom_objects={'CompositeLoss': CompositeLoss, 'CustomLoss': CustomLoss})

备选方案：函数式损失的临时修复（不推荐长期使用） #

如果不想修改为类实现，可以提前生成损失函数实例并注册具体的函数对象：

# 先生成具体的损失函数实例
alpha = 0.5
my_composite_loss = composite_loss(alpha)

# 注册这个实例（用固定名称）
utils.get_custom_objects()['my_composite_loss'] = my_composite_loss

# 编译时使用该实例
model.compile(loss=my_composite_loss, optimizer=optimizer)

# 保存模型
model.save('SYM-Revised-CAS-30')

# 加载时必须注册同一个名称对应的实例
model = load_model('SYM-Revised-CAS-30', custom_objects={'my_composite_loss': my_composite_loss})

注意：这种方法依赖于保存和加载时使用同一个函数实例，重启代码后需要重新生成实例并注册，稳定性不如类实现。

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