python - 如何替换 infs 以避免 PyTorch 中的 nan 梯度

Question

我需要计算 log(1 + exp(x)) 然后对其使用自动微分。但是对于太大的 x，它输出 inf 因为求幂:

>>> x = torch.tensor([0., 1., 100.], requires_grad=True)
>>> x.exp().log1p()
tensor([0.6931, 1.3133,    inf], grad_fn=<Log1PBackward>)

因为 log(1 + exp(x)) ≈ x 对于大的 x，我想我可以用 替换 infs >x 使用 torch.where。但是当这样做时，我仍然得到 nan 太大值的梯度。您知道为什么会发生这种情况吗？是否有其他方法可以让它发挥作用？

>>> exp = x.exp()
>>> y = x.where(torch.isinf(exp), exp.log1p())  # Replace infs with x
>>> y  # No infs
tensor([  0.6931,   1.3133, 100.0000], grad_fn=<SWhereBackward>)
>>> y.sum().backward()  # Automatic differentiation
>>> x.grad  # Why is there a nan and how can I get rid of it?
tensor([0.5000, 0.7311,    nan])

Answer 1

我发现的一个解决方法是手动实现 Log1PlusExp 函数及其向后的对应函数。然而，它并没有解释问题中 torch.where 的不良行为。

>>> class Log1PlusExp(torch.autograd.Function):
...     """Implementation of x ↦ log(1 + exp(x))."""
...     @staticmethod
...     def forward(ctx, x):
...         exp = x.exp()
...         ctx.save_for_backward(x)
...         return x.where(torch.isinf(exp), exp.log1p())
...     @staticmethod
...     def backward(ctx, grad_output):
...         x, = ctx.saved_tensors
...         return grad_output / (1 + (-x).exp())
... 
>>> log_1_plus_exp = Log1PlusExp.apply
>>> x = torch.tensor([0., 1., 100.], requires_grad=True)
>>> log_1_plus_exp(x)  # No infs
tensor([  0.6931,   1.3133, 100.0000], grad_fn=<Log1PlusExpBackward>)
>>> log_1_plus_exp(x).sum().backward()
>>> x.grad  # And no nans!
tensor([0.5000, 0.7311, 1.0000])