Keras/TensorFlow：run_eagerly设为None时自定义损失未生效？

我在Keras中通过以下方式添加自定义损失层：

custom_loss_layer = CustomLossLayer(params)(inputs, outputs)

model.add_loss(custom_loss_layer)

model.compile(..., loss="categorical_crossentropy", run_eagerly=True)

当设置run_eagerly=True并调用model.evaluate()时，总损失约为10+，符合预期（包含自定义损失，高于输出的联合损失）。但启用run_eagerly会导致严重性能问题。

若将run_eagerly设为None（默认图模式）再次执行evaluate()，总损失仅约为3，恰好等于所有输出的联合损失，自定义损失仿佛未被计算。

我怀疑问题出在自定义损失层的实现上：其中通过tf.squeeze(tf.argmax(input))将独热编码输入转为标量后，又转成普通整数用于访问列表索引。这种将张量转为原始值的操作，可能导致TensorFlow在图模式下无法追踪张量修改，进而自定义损失未被正确计算。需要说明的是，eager模式下损失函数工作正常，训练几轮后自定义损失可被正确最小化至约2（不计categorical_crossentropy损失）。

自定义损失层代码如下：

class IntervalPenaltyLayer(keras.layers.Layer):
def __init__(self, intelligible_intervals, pitch_mapping, **kwargs):
    self._intelligible_intervals = intelligible_intervals
    self._pitch_mapping = pitch_mapping
    super().__init__(**kwargs)

def call(self, inputs, outputs):
    factor = 5
    loss = 0

    input_current_tone_batch = inputs["tone_0"][:, -1, :]
    batch_size = input_current_tone_batch.shape[0]
    if batch_size is None:  # Check if it is still compiling
        return 0

    input_next_tone_batch = inputs["tone_1"][:, -1, :]
    input_pitch_batch = inputs["pitch"][:, -1, :]
    output_pitch_batch = outputs["pitch"]  # Output no need

    for i in range(batch_size):
        current_tone_tensor = input_current_tone_batch[i, :]
        next_tone_tensor = input_next_tone_batch[i, :]
        current_pitch_tensor = input_pitch_batch[i, :]
        predicted_pitch_tensor = output_pitch_batch[i, :]

        current_tone_idx = int(tf.squeeze(tf.argmax(current_tone_tensor, axis=-1)))
        next_tone_idx = int(tf.squeeze(tf.argmax(next_tone_tensor, axis=-1)))
        current_pitch_idx = int(tf.squeeze(tf.argmax(current_pitch_tensor, axis=-1)))

        current_pitch = self._pitch_mapping[current_pitch_idx]

        key = str(current_tone_idx) + "_" + str(next_tone_idx)
        if key not in self._intelligible_intervals:  # Rest or whatever, skipping
            continue

        possible_intervals = self._intelligible_intervals[key]

        for j in tf.range(tf.shape(predicted_pitch_tensor)[0]):
            new_pitch = self._pitch_mapping[int(j)]
            interval = new_pitch - current_pitch
            probability = tf.gather(predicted_pitch_tensor, j)
            if interval not in possible_intervals:
                loss += factor * probability

    return math.log(1 + loss)

你的怀疑完全正确，图模式下自定义损失失效的核心原因是代码中大量使用Python原生操作而非TensorFlow图兼容API，导致TensorFlow无法构建有效的计算图，损失项被直接忽略：

• int()强制转换张量：tf.argmax返回的是动态张量，用int()转为Python整数会丢失张量的追踪信息，图模式下TensorFlow无法识别这个值的来源和依赖关系，相当于把动态计算变成静态常量，无法纳入计算图。
• Python原生循环遍历批次：for i in range(batch_size)是Python循环，图模式下batch_size在编译阶段是未知的，循环内的损失计算完全不会被转化为图节点，相当于没执行。
• 字典/集合的原生查找：self._pitch_mapping[current_pitch_idx]、interval not in possible_intervals这类操作是Python原生逻辑，无法被TensorFlow图追踪，导致损失计算的依赖链断裂。
• math.log替代tf.math.log：math.log只能处理Python标量，无法兼容TensorFlow张量，图模式下会直接返回无效值（这里因为前面的问题已经让损失为0，所以返回math.log(1+0)=0）。

要让自定义损失在图模式下正常工作，必须把所有操作替换为TensorFlow的图兼容API，彻底抛弃Python原生的类型转换、循环和字典操作：

1. 将静态映射转为TensorFlow可处理的结构 #

• 把_pitch_mapping转为tf.constant，用tf.gather实现张量索引；
• 用tf.lookup.StaticHashTable把_intelligible_intervals转为图兼容的字典，避免原生字符串键查找。

2. 用向量化操作替代Python循环 #

去掉for i in range(batch_size)这类循环，改用tf.map_fn或向量化逻辑处理整个批次，确保所有计算都能被TensorFlow追踪。

3. 替换所有Python原生操作 #

• 去掉int()转换，直接用张量索引；
• 用tf.math.log替代math.log；
• 用tf.reduce_any实现音程是否允许的判断，替代Python的in操作。

修复后的示例代码（核心逻辑调整） #

class IntervalPenaltyLayer(keras.layers.Layer):
    def __init__(self, intelligible_intervals, pitch_mapping, **kwargs):
        # 将pitch_mapping转为TensorFlow常量
        self._pitch_mapping = tf.constant(pitch_mapping, dtype=tf.float32)
        # 构建音程查找表：将字符串键转为整数编码，值为允许的音程张量
        keys = []
        values = []
        max_tone = max(int(k.split("_")[0]) for k in intelligible_intervals.keys()) + 1
        for key, intervals in intelligible_intervals.items():
            current, next_ = map(int, key.split("_"))
            int_key = current * max_tone + next_
            keys.append(int_key)
            values.append(tf.constant(intervals, dtype=tf.float32))
        # 静态哈希表，默认返回空张量表示无允许音程
        self._interval_table = tf.lookup.StaticHashTable(
            tf.lookup.KeyValueTensorInitializer(keys, values, key_dtype=tf.int32, value_dtype=tf.float32),
            default_value=tf.constant([], dtype=tf.float32)
        )
        self._max_tone = max_tone
        self.factor = 5.0
        super().__init__(**kwargs)

    def call(self, inputs, outputs):
        # 提取批次数据，无需单独处理batch_size（图模式下动态适配）
        input_current_tone = inputs["tone_0"][:, -1, :]
        input_next_tone = inputs["tone_1"][:, -1, :]
        input_pitch = inputs["pitch"][:, -1, :]
        output_pitch = outputs["pitch"]

        # 获取批次内所有样本的索引（张量形式，无int转换）
        current_tone_idx = tf.argmax(input_current_tone, axis=-1, output_type=tf.int32)
        next_tone_idx = tf.argmax(input_next_tone, axis=-1, output_type=tf.int32)
        current_pitch_idx = tf.argmax(input_pitch, axis=-1, output_type=tf.int32)

        # 向量化获取当前音高
        current_pitch = tf.gather(self._pitch_mapping, current_pitch_idx)

        # 计算整数键，查找允许的音程
        interval_keys = current_tone_idx * self._max_tone + next_tone_idx
        allowed_intervals = self._interval_table.lookup(interval_keys)

        # 定义单个样本的损失计算逻辑，用tf.map_fn处理整个批次
        def compute_sample_loss(args):
            curr_pitch, allowed_intv, pred_pitch = args
            # 计算所有预测音高对应的音程
            new_pitches = self._pitch_mapping
            intervals = new_pitches - curr_pitch
            # 判断每个音程是否在允许列表中
            is_allowed = tf.reduce_any(tf.equal(tf.expand_dims(allowed_intv, 1), tf.expand_dims(intervals, 0)), axis=0)
            # 累加不允许音高的惩罚项
            penalty = self.factor * pred_pitch * tf.cast(tf.logical_not(is_allowed), tf.float32)
            return tf.reduce_sum(penalty)

        # 对整个批次应用损失计算
        sample_losses = tf.map_fn(
            compute_sample_loss,
            elems=(current_pitch, allowed_intervals, output_pitch),
            dtype=tf.float32
        )

        # 计算总损失，用tf.math.log替代math.log
        total_loss = tf.math.log(1.0 + tf.reduce_sum(sample_losses))
        return total_loss

额外注意事项 #

• 编译模型时无需设置run_eagerly=True，图模式下即可正常计算自定义损失；
• 如果intelligible_intervals的键是字符串，也可以直接用tf.lookup.StaticHashTable处理字符串键，无需转为整数编码；
• 向量化操作相比Python循环会大幅提升训练效率，这也是图模式的核心优势。

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