Keras迁移学习自定义损失调用tf.tensor_scatter_nd_update报错解决方案

需求描述 #

实现自定义二元交叉熵损失函数，当标签y_true取值为10**6时，将对应位置的标签和预测值都置为0，使这部分无标签样本不参与损失计算，不对模型参数更新产生贡献。

初始实现 #

def custom_loss_match(y_true, y_pred):
  '''自定义损失，当y_true等于10**6（代表样本无标签）时，对应位置的y_true和y_pred取值保持一致，从而忽略这部分样本的损失'''

  bool_mask = tf.equal(y_true, 10**6)
  indices = tf.where(bool_mask)
  if len(indices) != 0:
    y_true = tf.tensor_scatter_nd_update(y_true, indices, tf.cast([0]*len(indices), tf.int64))
    y_pred = tf.tensor_scatter_nd_update(y_pred, indices, tf.cast([0]*len(indices), tf.float32))

  return keras.losses.BinaryCrossentropy(from_logits=False)(y_true, y_pred)

问题现象 #

本地测试正常 #

使用自行构造的小批量样本测试可以得到合理结果：

custom_loss_match(tf.constant([[1],[0],[1],[0],[1],[0],[1],[10**6]], dtype=tf.int64), 
                  np.array([[np.random.random()] for _ in range(8)]))

模型训练报错 #

实际训练时抛出如下错误：

InvalidArgumentError:  Indices and updates specified for empty output shape
     [[{{node custom_loss_match/cond/then/_0/custom_loss_match/cond/TensorScatterUpdate}}]] [Op:__inference_train_function_326910]

Function call stack:
train_function

报错原因 #

问题出在tf.tensor_scatter_nd_update的调用逻辑与TensorFlow的静态图执行模式不兼容：

• 训练时TensorFlow会先构建静态计算图，len(indices)是Python层面的判断，无法在图构建阶段获取动态运行时的张量实际值，分支判断逻辑失效
• 当某个batch中不存在标签为10**6的样本时，indices是空张量，此时tf.tensor_scatter_nd_update拿到空的更新值列表，就会触发输出形状不匹配的错误

可行解决方案 #

使用掩码逐元素乘法替代条件判断和张量更新操作，所有运算均为张量层面的原生操作，完全兼容静态图执行模式：

def custom_loss_match(y_true, y_pred):
  '''自定义损失，当y_true等于10**6（代表样本无标签）时，对应位置的y_true和y_pred取值保持一致，从而忽略这部分样本的损失'''

  bool_mask = tf.not_equal(y_true, 10**6)

  y_true *= tf.cast(bool_mask, dtype=tf.int64)
  y_pred *= tf.cast(bool_mask, dtype=tf.float32)

  return keras.losses.BinaryCrossentropy(from_logits=False)(y_true, y_pred)

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