Help regarding debugging custom optimizer?

seankala: <p>Hi everyone. I’m trying to implement the elastic averaging stochastic gradient descent (EASGD) algorithm from the paper <a href="https://arxiv.org/abs/1412.6651" rel="noopener nofollow ugc"><em>Deep Learning with Elastic Averaging SGD</em></a> and was running into some trouble.</p>
<p>I’m using PyTorch’s <a href="https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer" rel="noopener nofollow ugc"><code>torch.optim.Optimizer</code></a> class and referencing the <a href="https://pytorch.org/docs/stable/_modules/torch/optim/sgd.html#SGD" rel="noopener nofollow ugc">official implementation of SGD</a> and the <a href="https://github.com/rahulkidambi/AccSGD/blob/master/AccSGD.py" rel="noopener nofollow ugc">official implementation of Accelerated SGD</a> in order to start off somewhere.</p>
<p>The code that I have is:</p>
<pre><code class="lang-auto">import torch.optim as optim


class EASGD(optim.Optimizer):
    def __init__(self, params, lr, tau, alpha=0.001):
        self.alpha = alpha

        if lr &lt; 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
</code></pre>
<p>When I run this code, I get the following error:</p>
<pre><code class="lang-auto">/home/user/github/test-repo/easgd.py:50: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the gradient for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations.
</code></pre>
<p>Reading <a href="https://discuss.pytorch.org/t/leaf-variable-was-used-in-an-inplace-operation/308">this PyTorch Discussion</a> helped understand what the difference between leaf and non-leaf variables are, but I’m not sure how I should fix my code to make it work properly.</p>
<p>Any tips on what to do or where to look are appreciated. Thanks.</p>

ptrblck: <p>This line of code:</p>
<pre><code class="lang-python">p = p - self.alpha * (x_normal - x_tilde)
</code></pre>
<p>would create a non-leaf tensor, so that the following <code>p.grad</code> call would raise this warning.<br>
Based on your code I think you should new a new variable name for this assignment and keep <code>p</code> pointing to the parameters in the optimizer.<br>
Also, don’t use the <code>.data</code> attribute, as it might have unwanted side effects and wrap the code in a <code>torch.no_grad()</code> block, if necessary.</p>