Save you from anything

0%

如何在keras中实现复杂的损失函数

发表于 分类于 阅读次数： Valine：

keras框架要求给定的损失函数有且只有y_pre和y_true两个入参。但许多损失函数并不是这个格式，比如推荐系统中用到的三元组损失，又比如变分词嵌入中变分推断模型的损失函数。因而需要用一些“奇技淫巧”才能在keras中实现复杂损失函数。

背景

在实现变分词嵌入版的itag时，我遇到了复杂损失函数的问题，使用了变分词嵌入的模型有两个损失函数，一个是主任务（标签推荐）的损失函数，一个是变分推断的损失函数。前者有对应的y_pre和y_true，而后者并没有，不仅没有，损失函数的运算步骤还非常复杂。

这导致了我无法使用keras框架，但itag本身是keras实现的，如果搬运到pytorch中，会耗费我大量的时间重构代码，最后我选择在keras中使用一些hook的方法实现复杂损失函数。

苏剑林有一篇利用embedding层实现复杂损失函数的博文，因为embedding本质上是一个矩阵，换句话说它可以当一个全连接层用，可以通过embedding实现一些复杂的损失函数，有兴趣的可以去看卡。但我用的不是这个方法。

问题分析

使用了变分词嵌入的模型有两个损失函数，一个是主任务的损失一个是变分词嵌入的损失，其中变分词嵌入的损失的具体实现如下（有修改）：

1
2
3
4
5
6
7
8
9
...
weight_decay, learning_rate = get_decay_rate(epochs)
...
reg_loss = weight_decay * embedding.regularizer()
...
loss = cross_entropy + reg_loss
...
self.optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate).minimize(loss, global_step=global_step, var_list=trainables)
...
  1. 根据当前的epoch值生成权重递减值和学习率（具体算法不重要）
  2. 调用嵌入层中的一个方法生成变分损失
  3. 对这个损失进行权重递减
  4. 和主任务损失相加
  5. 反向传播，用了adam优化器但是仍然每轮手动调整学习率

上面的embedding.regularizer()方法实现如下：

1
2
3
4
5
6
# KL散度（变分估计损失）计算函数
def regularizer(self):
    k1, k2, k3 = 0.63576, 1.8732, 1.48695
    log_alpha = self.clip(self.embedding_logdropout_ratio) 
    KLD = -tf.reduce_sum(k1 * tf.sigmoid(k2 + k3 * log_alpha) - 0.5 * tf.nn.softplus(-log_alpha) - k1)
    return KLD

即变分词嵌入的损失的计算实际上需要的外部输入只有一个epoch，并不符合keras框架的对损失函数要求，因此不能使用keras的loss设定方法。另外还有一个问题，在keras框架中，epoch中无法从被训练的模型内部获取，因而也需要一个解决方案。

对keras进行hooc

通过定义一个“损失层”，在层内完成复杂损失的计算逻辑：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class ComplexLoss(Layer):
    def __init__(self, **kwargs):
        super(ComplexLoss, self).__init__(**kwargs)
        
    def build(self, input_shape):
        super(ComplexLoss, self).build(input_shape)

    def call(self, inputs,  **kwargs):
        # 由于keras限制Layer的子类必须以张量或张量列表入参，这里需要手动解参。
        '''
        y_true, y_pred用于计算主任务损失，即y_true也需要作为输入数据。
        由之前计算KL散度的代码可知，计算KL散度需要外部输入epoch和变分嵌入层的embedding_logdropout_ratio。
        为了代码的解耦（不在嵌入层内进行损失的计算），由嵌入层将embedding_logdropout_ratio作为返回值，传递给ComplexLoss层。
        '''
        y_true, y_pred, embedding_logdropout_ratio, epoch = inputs
        # 计算主任务损失
        cross_loss = mean_negative_log_probs(y_true, y_pred)
        # 计算变分嵌入损失，包括KL散度、权重递减等所有操作都在这一个方法内完成
        reg_loss = reg_loss_function(embedding_logdropout_ratio, epoch)
        # 获得总损失
        main_loss = cross_loss + reg_loss
        # 手动向Layer添加损失（以下两行为核心代码）
        self.add_loss(main_loss, inputs=True)
        # 添加到metric中，就能在训练的时候看到keras打印这个值了，否则keras只是会拿main_loss进行反向传播，而不会在训练时打印在实时进度中。
        self.add_metric(main_loss, aggregation="mean", name="loss")
        # 使用自制损失层时，这个层会在模型的最后，需要返回值以完成keras框架的构建
        return main_loss

在使用一个定制loss层后，loss的计算在这个层内完成了，因而keras的外围代码也需要修改：

（以下代码为简化版，仅为说明ComplexLoss的用法，实际上无法运行）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
# 创建模型的各个层
encoder_input = Input()
decoder_input = Input()
# 用于ComplexLoss层的y_true
y_true_input = Input()
# 由于模型无法从keras框架处获得epoch值，这里通过input层向模型传递当前的epoch值。
# epoch的更新逻辑写在Data generator中，每循环一轮数据epoch+1。
epoch_input = Input()

# 搭建seq2seq模型
encoder = EnCoder()
decoder = Decoder()
loss_layer = ComplexLoss()
decoder_dense = Dense()

# 解码器内嵌在encoder层，所以由encoder返回embedding_logdropout_ratio
encoder_outputs, state， embedding_logdropout_ratio = encoder(encoder_input)
decoder_outputs, n_state = decoder(encoder_outputs, initial_state=state)
output = decoder_dense(decoder_outputs)
main_loss = loss_layer([y_true_input, output, embedding_logdropout_ratio, epoch_input])
# 用main_loss作为模型的最终输出，这个数据没有用了，只是为了完成模型的搭建
model = Model(inputs=[encoder_input, 
                      decoder_input, 
                      y_true_input, 
                      epoch_input], 
              outputs=[main_loss])

adam = optimizers.Adam(lr=0.001)
# 这里不再需要设定loss
model.compile(optimizer=adam)
model.fit()

以上代码和普通的keras代码的两个主要区别是：

  • y_true需要作为input的一部分（以传递给损失层）
  • 不再需要在compile中指定损失函数（因为已经通过损失层实现了）

总结

使用以上损失层的方法，几乎可以在keras中定义任何形式的损失函数而不受框架的限制。