具有多个输入的Tensorflow 2.0自定义损失功能

Question

我正在尝试使用以下两个损失函数来优化模型

def loss_1(pred, weights, logits):
    weighted_sparse_ce = kls.SparseCategoricalCrossentropy(from_logits=True)
    policy_loss = weighted_sparse_ce(pred, logits, sample_weight=advantages)

和

def loss_2(y_pred, y):
    return kls.mean_squared_error(y_pred, y)

但是，由于TensorFlow 2期望损失函数的形式为

def fn(y_pred, y_true):
    ...

我正在使用 loss_1的变通方法，在将 pred 的调用中传递给 weights之前，将 loss_1和 model.fit打包到单个张量中，然后将它们解压缩到 loss_1中。这是优雅且令人讨厌的，因为 pred和 weights是不同的数据类型，因此每次我调用 model.fit 时，都需要进行额外的强制转换，打包，解压缩和取消转换。

此外，我知道 sample_weight的 fit参数，有点像 this question的解决方案。如果不是因为我正在使用两个损失函数，而我只希望将 sample_weight应用于其中一个函数，则这可能是一个可行的解决方案。同样，即使这是一个解决方案，也无法将其推广到其他类型的自定义损失函数。

简而言之，我的问题是:

创建具有任意数量的损失函数的最佳方法是什么
TensorFlow 2中的参数？

我尝试过的另一件事是传递 tf.tuple ，但这似乎也违反了TensorFlow对损失函数输入的期望。

Answer 1

使用TF2中的custom training可以轻松解决此问题。您只需要在GradientTape上下文中计算两分量损失函数，然后使用产生的梯度调用优化器即可。例如，您可以创建一个函数custom_loss，计算给定每个参数的两个损耗:

def custom_loss(model, loss1_args, loss2_args):
  # model: tf.model.Keras
  # loss1_args: arguments to loss_1, as tuple.
  # loss2_args: arguments to loss_2, as tuple.
  with tf.GradientTape() as tape:
    l1_value = loss_1(*loss1_args)
    l2_value = loss_2(*loss2_args)
    loss_value = [l1_value, l2_value]
  return loss_value, tape.gradient(loss_value, model.trainable_variables)

# In training loop:
loss_values, grads = custom_loss(model, loss1_args, loss2_args)
optimizer.apply_gradients(zip(grads, model.trainable_variables))

这样，每个损失函数都可以采用任意数量的渴望张量，无论它们是模型的输入还是输出。每个损失函数的参数集不需要是不相交的，如本示例所示。