如何在TensorFlow 2 Keras中打印model.add_loss()生成的损失值

问题原因 #

通过model.add_loss()添加的损失会被Keras自动纳入总损失计算，但默认情况下，训练过程中只会展示model.compile()时指定的loss和metrics，不会单独输出add_loss()添加的损失项。

解决方案 #

方法1：自定义Metric单独监控KL散度损失 #

自定义一个继承tf.keras.metrics.Metric的类，用来计算并记录KL散度损失，在compile时将其加入metrics列表，训练时就会自动打印该损失值。

修改后的完整代码：

import tensorflow as tf
keras = tf.keras
import keras.backend as K
import numpy as np

def ohc(y):
    y = np.random.randint(4, size=100)
    b = np.zeros((y.size, y.max() + 1))
    b[np.arange(y.size), y] = 1
    return b

xtrain = np.random.rand(100,50)
y = np.random.randint(4, size=100)
y_one_hot = ohc(y)

# 自定义KL散度监控Metric
class KLDivergenceMetric(keras.metrics.Metric):
    def __init__(self, name='kl_divergence', **kwargs):
        super().__init__(name=name, **kwargs)
        self.kl_loss = self.add_weight(name='kl', initializer='zeros')
        self.count = self.add_weight(name='count', initializer='zeros')

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_pred2 = y_pred + 0.00001
        LR = y_true / y_pred2
        logLR = K.log(LR)
        kld = y_true * logLR
        loss = K.mean(kld) * 0.2
        
        self.kl_loss.assign_add(loss)
        self.count.assign_add(1.)

    def result(self):
        return self.kl_loss / self.count

def my_kld(y_true, y_pred):
    y_pred2 = y_pred + 0.00001
    LR = y_true / y_pred2
    logLR = K.log(LR)
    kld = y_true * logLR
    loss = K.mean(kld) * 0.2
    return loss

x_t = keras.layers.Input((50,))
y_t = keras.layers.Input((4,))

x1 = keras.layers.Dense(100, activation='relu')(x_t)
x2 = keras.layers.Dense(4, activation='softmax')(x1)

model = keras.models.Model([x_t,y_t], x2)
model.add_loss(my_kld(y_t, x2))

optim = keras.optimizers.Nadam(0.00006) 
# 将自定义Metric加入监控列表
model.compile(loss=['categorical_crossentropy'], optimizer=optim, metrics=['accuracy', KLDivergenceMetric()])
model.fit(x=[xtrain,y_one_hot], y=y_one_hot, epochs=10)

方法2：用回调函数手动打印损失 #

如果不需要长期监控，可通过自定义Callback在每个epoch结束时，从模型的losses属性中提取add_loss()添加的损失值并打印。

示例代码（仅修改compile和fit部分）：

# 自定义回调函数
class PrintKLLoss(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        # model.losses包含所有add_loss添加的损失，取第一个即为KL散度损失
        kl_loss = K.eval(self.model.losses[0])
        print(f"\nEpoch {epoch+1}: KL Divergence Loss = {kl_loss:.4f}")

optim = keras.optimizers.Nadam(0.00006) 
model.compile(loss=['categorical_crossentropy'], optimizer=optim, metrics=['accuracy'])
# 加入回调函数
model.fit(x=[xtrain,y_one_hot], y=y_one_hot, epochs=10, callbacks=[PrintKLLoss()])

方法3：将KL散度作为模型输出之一 #

修改模型结构，把KL散度损失作为额外输出，在compile时为每个输出指定损失规则，训练时会分别显示各项损失。

示例代码（修改模型定义和compile部分）：

# 修改模型结构，添加KL损失输出
x_t = keras.layers.Input((50,))
y_t = keras.layers.Input((4,))

x1 = keras.layers.Dense(100, activation='relu')(x_t)
x2 = keras.layers.Dense(4, activation='softmax')(x1)

kl_loss = my_kld(y_t, x2)
# 模型包含两个输出：分类结果、KL损失
model = keras.models.Model([x_t,y_t], [x2, kl_loss])

optim = keras.optimizers.Nadam(0.00006) 
# 为每个输出指定损失，KL损失直接返回自身即可
model.compile(
    loss=['categorical_crossentropy', lambda y_true, y_pred: y_pred],
    optimizer=optim,
    metrics=['accuracy']
)
# fit时传入分类标签和dummy的0数组作为KL损失的目标值
model.fit(
    x=[xtrain,y_one_hot], 
    y=[y_one_hot, np.zeros((100,))], 
    epochs=10
)

说明 #

• 方法1适合长期监控，损失值会和准确率一样自动出现在训练日志中；
• 方法2适合临时调试，快速查看损失数值；
• 方法3需要调整模型结构，适合需要将损失作为输出的场景。

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