pytorch - 如何在 PyTorch 中计算幂函数 wrt 指数的梯度？

Question

我正在尝试计算的梯度

out = x.sign()*torch.pow(x.abs(), alpha)

关于阿尔法。

到目前为止，我尝试了以下方法：

class Power(nn.Module):
  def __init__(self, alpha=2.):
    super(Power, self).__init__()
    self.alpha = nn.Parameter(torch.tensor(alpha))

  def forward(self, x):
    return x.sign()*torch.abs(x)**self.alpha

但是这门课一直nan在训练我的网络。我希望看到类似的东西，grad=out*torch.log(x)但无法做到。例如，此代码不返回任何内容：

alpha_rooting = Power()
x = torch.randn((1), device='cpu', dtype=torch.float)
out = (alpha_rooting(x)).sum()
out.backward()
print(out.grad)

我autograd也没有运气尝试使用它。我应该如何解决这个问题？谢谢。

Answer 1

您编写的Power()课程按预期工作。您实际使用它的方式存在问题。梯度存储在.grad该变量中，而不是out您在上面使用的变量中。您可以更改代码如下。

alpha_rooting = Power()
x = torch.randn((1), device='cpu', dtype=torch.float)
out = (alpha_rooting(x)).sum()

# compute gradients of all parameters with respect to out (dout/dparam)
out.backward()
# print gradient of alpha
# Note that gradients are store in .grad of parameter not out variable
print(alpha_rooting.alpha.grad)

# compare if it is approximately correct to exact grads
err = (alpha_rooting.alpha.grad - out*torch.log(x))**2 
if (err <1e-8):
    print("Gradients are correct")