What is the correct way of computing a grad penalty using AMP?

username: ayalaa2
cooked: <p>I’m converting my distributed code to work with PyTorch’s AMP, but I have confusion on how I should compute my grad penalty as part of my WGAN-GP loss. I have referred to the <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-penalty" rel="nofollow noopener">documentation</a> on this matter but I believe my use case is slightly different.</p>
<p>Here’s my code on computing the grad penalty:</p>
<pre><code class="lang-auto">out = discrim(sample)

gradients = torch.autograd.grad(inputs=sample, outputs=out,
    grad_outputs=torch.ones(out.shape).to(sample.device),
    create_graph=True, retain_graph=True, only_inputs=True)[0]
</code></pre>
<p>You may assume that sample is a tensor yielded by my generator and out is a <code>N x 1</code> tensor. I’ll start off with saying that this code <em>works</em> without amp, but I’m just confused on how I’m suppose to use amp with this. This is what I’m currently trying:</p>
<pre><code class="lang-auto">with autocast(False):
    sample = scaler.scale(sample)
    out = discrim(sample)

    gradients = torch.autograd.grad(inputs=sample, outputs=out,
        grad_outputs=torch.ones(out.shape).to(sample.device),
        create_graph=True, retain_graph=True, only_inputs=True)[0]
    gradients = gradients / scaler.get_scale()
</code></pre>
<p>Note that <code>sample</code> was computed under <code>autocast(True)</code>, but I’m not sure if this matters. This seems to work, but I think I want to compute <code>out</code> under the mixed precision context to be faster. This seems to be slow atm.</p>
<p>This my second attempt at the problem:</p>
<pre><code class="lang-auto">with autocast(True):
    out = discrim(x=sample)

    with autocast(False):
        sample = scaler.scale(sample)
        out = scaler.scale(out)

        gradients = torch.autograd.grad(inputs=sample, outputs=out,
            grad_outputs=torch.ones(out.shape).to(sample.device),
            create_graph=True, retain_graph=True, only_inputs=True)[0]

        gradients = gradients / scaler.get_scale()
</code></pre>
<p>But I get this error on autograd: <code>One of the differentiated Tensors appears to not have been used in the graph</code>. I’m not really sure on what’s the proper way of doing this. Any help is greatly appreciated!</p>

username: mcarilli
cooked: <p>I don’t believe your use case is different from what’s shown in <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-penalty" rel="noopener nofollow ugc">the docs</a>, aside from more explicit kwargs and <code>outputs</code> being a Tensor with (presumably) more than one element.  Perhaps, instead of</p>
<blockquote>
<p>To implement a gradient penalty <em>with</em> gradient scaling, the loss passed to <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.grad" rel="noopener nofollow ugc"> <code>torch.autograd.grad()</code> </a> should be scaled.</p>
</blockquote>
<p>the docs should say</p>
<blockquote>
<p>To implement a gradient penalty <em>with</em> gradient scaling, the <strong><code>outputs</code> Tensor(s)</strong> passed to <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.grad" rel="noopener nofollow ugc"> <code>torch.autograd.grad()</code> </a> should be scaled.</p>
</blockquote>
<p>Try</p>
<pre><code class="lang-auto">with autocast():
    out = discrim(sample)

gradients = torch.autograd.grad(outputs=scaler.scale(out), inputs=sample,
    grad_outputs=torch.ones(out.shape, device=sample.device),
    create_graph=True, retain_graph=True, only_inputs=True)[0]

# proceed as shown in the docs

# if gradients is a Tensor
gradients = gradients / scaler.get_scale()
# if gradients is a tuple
inv_scale = 1./scaler.get_scale()
gradients = [p * inv_scale for p in gradients]

with autocast():
    # compute penalty term from gradients
    # add to loss if needed

scaler.scale(&lt;loss or penalty&gt;).backward()
</code></pre>
<p>If that works I will update docs as described above.</p>
<p>FYI in <code>grad</code>, <code>retain_graph</code> defaults to the value of <code>create_graph</code> and <code>only_inputs</code> defaults to <code>True</code> so these kwargs are not needed.  Also, I changed<br>
<code>torch.ones(out.shape).to(sample.device)</code> to <code>torch.ones(out.shape, device=sample.device)</code> which synthesizes <code>grad_outputs</code> directly on the device, avoids an expensive CPU-&gt;GPU copy, and is good practice regardless of Amp.</p>

username: ayalaa2
cooked: <p>Thank you for the reply!</p>
<p>I’ve implemented the suggestion you made and it all appears to be working. With your comment and cross-referencing with the docs, it makes much more sense on what’s happening now. Thank you also for the additional comments, I didn’t realize the <code>to(...)</code> had that problem and I made that change.</p>
<p>I do agree that if the docs listed it as the <strong>output tensors</strong> that this would make more sense for newcomers in the future.</p>