python - 如何计算 Tensorflow 中不变风险最小化的惩罚？

Question

我正在尝试实现一种称为“不变风险最小化”的技术，该技术在训练机器学习模型的损失函数中添加了一个惩罚项。新惩罚项的技术定义是关于常数分类器的平方梯度范数。这里有一个使用 PyTorch 的“惩罚”功能的实现。

我想知道如何在 Tensorflow 2 中实现这个功能。

更具体地说，我想实现下面的功能，这也在我分享它的链接的代码中。

  def penalty(logits, y):
    scale = torch.tensor(1.).cuda().requires_grad_()
    loss = mean_nll(logits * scale, y)
    grad = autograd.grad(loss, [scale], create_graph=True)[0]
    return torch.sum(grad**2)

标签： pythontensorflowmachine-learningloss-function

Answer 1

以类似的方式：

def penalty(y_true, y_pred):
    scale = tf.constant(1.)
    with tf.GradientTape() as tape:
        tape.watch(scale)
        loss = tf.losses.binary_crossentropy(y_true, y_pred*scale, from_logits=True)
    grad = tape.gradient(loss, [scale])[0]
    return tf.reduce_sum(grad**2)

请注意，与 PyTorch 版本相比，参数 logits 和 ground truth 的顺序是相反的，以尊重 TensorFlow 的约定。

要计算梯度，您只需要tf.GradientTape，您可以在指南中阅读更多内容：梯度和自动微分简介

---

比较两个版本产生相同的结果：

PyTorch 版本：

>>> penalty(torch.tensor(1.),torch.tensor(0.))
tensor(0.5344, grad_fn=<SumBackward0>)

>>> penalty(torch.tensor(1.),torch.tensor(1.))
tensor(0.0723, grad_fn=<SumBackward0>)

TensorFlow 版本：

>>> penalty([0.],[1.])
<tf.Tensor: shape=(), dtype=float32, numpy=0.53444666>

>>> penalty([1.],[1.])
<tf.Tensor: shape=(), dtype=float32, numpy=0.0723295>