tensorflow小技巧--时间序列模型中的复杂损失函数计算方法

BERT的MLM是我见过的比较复杂的优化任务了，我以此为例分享一下对于这种比较复杂的损失函数如何计算。先说下这个MLM是什么，如下图所示（摘自BERT的原始论文）：

所以这个任务就是将sentence中一些token进行掩盖，模型会输出这些掩盖的token的隐藏状态，将这些隐藏状态输入softmax可以得到候选单词的概率分布，这样根据ground truth就可以计算cross entropy了。原理比较简单，要想使用tf高效实现，博主着实费了一些脑细胞。我的计算方法如下：

import tensorflow as tf
final_output=tf.random.normal([3,4,8],
               name='final_output') #[batch_size,seq_len,category_size]
mask=tf.constant([[0,1,0,0],
                  [1,0,0,0],
                  [0,0,1,0]],dtype=tf.int32)
mask_value=tf.constant([[0,3,0,0],
                        [6,0,0,0],
                        [0,0,5,0]],dtype=tf.int32)

final_output就是模型最终的输入了，即softmax的输出结果，咱们假设类别共有8类。mask标记哪些token被掩盖了，掩盖的标1。mask_value是被掩盖的token的真实类别。我的想法就是先不管是否掩盖，计算所有token的cross entropy，然后只保留被掩盖的token的cross entropy，其余的cross entropy记为0，最后相加就平均就好啦。

mask=tf.cast(mask,dtype=tf.float32)
loss_all=tf.keras.losses.sparse_categorical_crossentropy(mask_value,final_output) #[batch_size,seq_len]
loss=tf.reduce_sum(loss_all*mask)/tf.reduce_sum(mask)

大家要是有更高效的计算方法，欢迎留言交流哈~~