如何在TensorFlow模型的不同epoch间动态切换损失函数

针对你每隔固定轮次更换损失函数的需求，以下是两种可行的修正方案，解决你之前遇到的变量更新失效、model.add_loss不生效的问题：

一、修复损失包装器的变量更新问题 #

你之前遇到current_epoch无法同步到包装器的核心原因是不可变变量的闭包作用域限制——整数属于不可变类型，回调里的更新无法同步到损失包装器的闭包环境中。改用可变对象（比如字典）存储轮次即可解决：

import tensorflow as tf
from tensorflow.keras.callbacks import Callback

# 用可变字典存储当前轮次，确保闭包能同步获取更新
epoch_tracker = {"current_epoch": 0}

def dynamic_loss(y_true, y_pred):
    if epoch_tracker["current_epoch"] % 5 == 0:
        # 第0、5、10...轮用损失函数A
        return tf.keras.losses.MSE(y_true, y_pred)
    else:
        # 其他轮用损失函数B
        return tf.keras.losses.BinaryCrossentropy()(y_true, y_pred)

class LossSwitchCallback(Callback):
    def on_epoch_end(self, epoch, logs=None):
        # 更新轮次（注意keras的epoch从0开始计数）
        epoch_tracker["current_epoch"] = epoch + 1

# 编译模型时使用包装后的动态损失
model.compile(optimizer="adam", loss=dynamic_loss)
# 添加轮次跟踪回调
model.fit(x_train, y_train, epochs=20, callbacks=[LossSwitchCallback()])

二、正确使用model.add_loss动态切换损失 #

之前model.add_loss不生效是因为两个关键点没做到：

1. 模型编译时已指定固定损失，add_loss添加的是额外损失而非替换原损失
2. 切换损失前未清除之前添加的自定义损失

修正方案如下：

class DynamicLossCallback(Callback):
    def __init__(self, loss_fn_a, loss_fn_b):
        self.loss_fn_a = loss_fn_a
        self.loss_fn_b = loss_fn_b

    def on_epoch_begin(self, epoch, logs=None):
        # 清除之前的自定义损失，避免叠加
        self.model.losses = []
        # 根据轮次选择目标损失函数
        if (epoch + 1) % 5 == 0:
            current_loss = self.loss_fn_a
        else:
            current_loss = self.loss_fn_b
        # 添加动态损失，绑定模型的输入输出
        def loss_wrapper():
            return current_loss(self.model.target[0], self.model.output[0])
        self.model.add_loss(loss_wrapper)

# 编译模型时不指定固定损失（设为None）
model.compile(optimizer="adam", loss=None)
# 定义两种待切换的损失函数
loss_a = tf.keras.losses.MSE
loss_b = tf.keras.losses.BinaryCrossentropy()
# 添加动态损失切换回调
model.fit(x_train, y_train, epochs=20, callbacks=[DynamicLossCallback(loss_a, loss_b)])

注意：如果你的模型是多输入/多输出结构，需要调整model.target和model.output的索引，匹配对应的数据维度。

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