自定义PyTorch优化器工作不正常，无法访问渐变

Question

我试图实现论文中的弹性平均随机梯度下降(EASGD(算法，但遇到了一些麻烦。

我使用PyTorch的torch.optim.Optimizer类，并参考SGD的官方实现和Accelerated SGD的正式实现，以便从某个地方开始。

我的代码是：

import torch.optim as optim

class EASGD(optim.Optimizer):
def __init__(self, params, lr, tau, alpha=0.001):
self.alpha = alpha
if lr < 0.0:
raise ValueError(f"Invalid learning rate {lr}.")
defaults = dict(lr=lr, alpha=alpha, tau=tau)
super(EASGD, self).__init__(params, defaults)
def __setstate__(self, state):
super(EASGD, self).__setstate__(state)
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
tau = group['tau']
for t, p in enumerate(group['params']):
x_normal = p.clone()
x_tilde = p.clone()
if p.grad is None:
continue
if t % tau == 0:
p = p - self.alpha * (x_normal - x_tilde)
x_tilde = x_tilde + self.alpha * (x_normal - x_tilde)
d_p = p.grad.data
p.data.add_(d_p, alpha=-group['lr'])
return loss

当我运行此代码时，我得到以下错误：

/home/user/github/test-repo/easgd.py:50:UserWarning:正在访问非叶张量的张量的.grad属性。在autograd.backward()期间不会填充其.grad属性。如果确实需要非叶张量的梯度，请在非叶张量上使用.retain_grad()。如果您错误地访问了非叶张量，请确保您访问的是叶张量。有关更多信息，请访问github.com/pytorch/pytarch/pull/30531。

阅读PyTorch讨论有助于理解叶变量和非叶变量之间的区别，但我不确定应该如何修复代码以使其正常工作。

任何关于该做什么或在哪里寻找的建议都将不胜感激。谢谢

Answer 1

我认为问题是您在以下行复制p：

p = p - self.alpha * (x_normal - x_tilde)

如果这一行被执行(第一个循环中t=0的情况也是如此(，则下一行将引发错误，因为p不再具有.grad属性。

您应该使用就地运算符add_、mult_、sub_、divide_等。

for t, p in enumerate(group['params']):
if p.grad is None:
continue
d_p = p.grad.data
if t % tau == 0:
d_p.sub_(self.alpha*0.01)
p.data.add_(d_p, alpha=-group['lr'])

上面，我删除了x_normal和x_tilde，因为您没有给它们适当的值。但我希望你明白。仅在处理step函数内部的数据时使用就地运算符。