Trying to calculate gradient penalty with grad() detaches computational graph

username: import-antigravity
cooked: <p>I have a very simple discriminator for a toy GAN problem where I’m trying to find the magnitude of the gradient in order to apply a penalty to the gradient. In order to do that, I need the gradient norm to be differentiable.</p>
<p>When I calculate the loss function for the generator I get the following computational graph:</p>
<pre><code class="lang-auto"># This code produces a tensor with a GradFn
bce = -1 * d_z.mean()
</code></pre>
<p>And then when I differentiate the loss with respect to the parameters I get a valid graph:</p>
<pre><code class="lang-auto"># This code produces a tensor with a GradFn
gradient = grad(bce, gan.G.parameters(), retain_graph=True, create_graph=True)
</code></pre>
<p>Here is the code for the generator:</p>
<pre><code class="lang-auto">class Linear(nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super(Linear, self).__init__()
        self.W = nn.Linear(in_features, out_features)

    def forward(self, x: Tensor) -&gt; Tensor:
        return self.W(x)
</code></pre>
<p>Now… when I do the same thing for the discriminator, it doesn’t work. The loss works fine:</p>
<pre><code class="lang-auto"># This code produces a tensor with a GradFn
d_x = gan.D(X)
d_z = gan.D(g_X.detach())
bce = d_z.mean() - d_x.mean()
</code></pre>
<p>But when I try to differentiate it again I get no graph:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/2/72b1e8b179aa159fafba28e14df8a1c4ef952b83.png" data-download-href="https://discuss.pytorch.org/uploads/default/72b1e8b179aa159fafba28e14df8a1c4ef952b83" title="Screen Shot 2020-10-20 at 8.51.22 PM"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/2/72b1e8b179aa159fafba28e14df8a1c4ef952b83.png" alt="Screen Shot 2020-10-20 at 8.51.22 PM" data-base62-sha1="gmDBbGViiEYlAIscOZOF4wBG67V" width="690" height="96" data-dominant-color="EAEDE8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2020-10-20 at 8.51.22 PM</span><span class="informations">698×98 11 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>And here is the code for the discriminator. As you can see, I’m only using building blocks from <code>torch.nn</code> so <code>requires_grad</code> should be true for everything:</p>
<pre><code class="lang-auto">class Quadratic(nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super(Quadratic, self).__init__()
        self.a = nn.Bilinear(in_features, in_features, out_features, bias=False)
        self.b = nn.Linear(in_features, out_features, bias=False)

    def forward(self, x: Tensor) -&gt; Tensor:
        return self.a(x, x) + self.b(x)
</code></pre>
<p>I’ve been trying to debug this all day, any help would be greatly appreciated!</p>

username: albanD
cooked: <p>Hi,</p>
<p>Do you set <code>requires_grad</code> afterwards for some reason on the params of gan.D ?<br>
Are you sure that your function is twice differentiable with respect to your parameters (and non-zero) though?</p>
<p>What happens if for the sake of the experiment you do bce = bce.exp() to ensure the function is properly differentiable?</p>

username: import-antigravity
cooked: <p>I’m using just nn.Linear() and nn.Bilinear() layers so it should be twice differentiable.</p>
<p>Well I’ll be damned, when I do that I get a proper graph:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/d/2dd72070b72b2397ddb0686dd0034f510ed6dcda.png" data-download-href="https://discuss.pytorch.org/uploads/default/2dd72070b72b2397ddb0686dd0034f510ed6dcda" title="Screen Shot 2020-10-21 at 11.28.15 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/d/2dd72070b72b2397ddb0686dd0034f510ed6dcda_2_343x500.png" alt="Screen Shot 2020-10-21 at 11.28.15 AM" data-base62-sha1="6xwomQTrk4cFjEqHD03JW0raNrI" width="343" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/d/2dd72070b72b2397ddb0686dd0034f510ed6dcda_2_343x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/d/2dd72070b72b2397ddb0686dd0034f510ed6dcda_2_514x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/d/2dd72070b72b2397ddb0686dd0034f510ed6dcda_2_686x1000.png 2x" data-dominant-color="F0F1F1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2020-10-21 at 11.28.15 AM</span><span class="informations">824×1198 75.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: albanD
cooked: <p>They are twice differentiable but at least for Linear, the second derivative is 0. And the autograd detects that the gradient is independent of the input and does not create graph for it.</p>

username: import-antigravity
cooked: <p>Well that’s the thing… The equation for the generator is Ax + b, so if anything it should be the one not returning a gradient, but it does.</p>

username: albanD
cooked: <p>Unfortunately in the autograd, we don’t make a difference between “is independent” and “is 0”.<br>
Also such gradient can be represented as:</p>
<ul>
<li>A Tensor full of 0s</li>
<li>None (or undefined tensor on the c++ side)</li>
<li>An error if you use <code>autograd.grad(..., allow_unused=False)</code>
</li>
</ul>
<p>And the reason for this is that sometimes it will create a graph that produces only 0s and sometimes won’t create the graph at all.<br>
The problem is that it is very hard to be always consistent here as we would like to never create the graph ideally, but we don’t want to have to do extra work either to know if we should create it or not.</p>

username: import-antigravity
cooked: <p>I feel like there might be a miscommunication here. Ultimately, what I’m trying to do is get the 2-norm of the gradient (first derivative), and add it as a term to the overall loss function. While it is true that the generator is linear, I’ve worked the math out in closed form and the second derivative (Hessian) of the <em>loss function</em> with respect to the parameters is in fact non-zero. Once again, the generator returns a proper graph and the gradient is non-zero, as expected. The problem is the discriminator, which is bilinear and should have a non-zero hessian no matter what. Sorry if I’m not doing a good job of explaining.</p>

username: albanD
cooked: <p>But the example you showed takes the mean, not the 2 norm! That makes a big difference!</p>

username: import-antigravity
cooked: <p>The mean is to aggregate the losses into a single value</p>

username: albanD
cooked: <p>Ho that is a different part of the code ok!<br>
The thing is that if the autograd does not create the graph, that means that the gradient will just always be 0 based on what you computed.<br>
Maybe you can share a full code sample that shows your problem?</p>

username: import-antigravity
cooked: <p>Here is a self-contained example, but as you can see it actually works now! This confirms that there is a bug somewhere in my code. I’ll take a closer look and report back once I find it:</p>
<pre><code class="lang-auto">import torch as tr
from torch.autograd import grad

# Generate data, 1000 samples from a normal distribution
X = tr.normal(tr.ones(1000)).view(-1, 1)
Z = tr.normal(tr.zeros(1000)).view(-1, 1)

# Define the generator and discriminator
# Generator: Az + b
class G(tr.nn.Module):
    def __init__(self):
        super(G, self).__init__()
        self.Ab = tr.nn.Linear(1, 1)
    
    def forward(self, x):
        return self.Ab(x)

# Discriminator: x^T C x + d^T x
class D(tr.nn.Module):
    def __init__(self):
        super(D, self).__init__()
        self.C = tr.nn.Bilinear(1, 1, 1, bias=False)
        self.d = tr.nn.Linear(1, 1, bias=False)
    
    def forward(self, x):
        return self.C(x, x) + self.d(x)

g = G()
d = D()

# Test to make sure outputs are valid
print("Generator Test:", g(Z).mean())
print("Discriminator Test:", d(X).mean(), d(g(Z)).mean())

# Define loss:
generator_loss = -1 * d(g(Z)).mean()
discriminator_loss = d(g(Z)).mean() - d(X).mean()

# Test to make sure loss is differentiable
print("Generator Loss:", generator_loss)
print("Discriminator Loss:", discriminator_loss)

# Get gradient wrt parameters
generator_gradient = grad(generator_loss, g.parameters(), retain_graph=True, create_graph=True)
discriminator_gradient = grad(discriminator_loss, d.parameters(), retain_graph=True, create_graph=True)

# Both tensors should still have grad_fn
print("Generator Gradient:", generator_gradient)
print("Discriminator Gradient:", discriminator_gradient)

# Take 2-norm of gradient components
generator_norm = tr.norm(tr.cat([tr.flatten(i) for i in generator_gradient]))
discriminator_norm = tr.norm(tr.cat([tr.flatten(i) for i in discriminator_gradient]))

# Add to loss function, should STILL have grad_fn
final_gen_loss = generator_loss + generator_norm
final_disc_loss = discriminator_loss + discriminator_norm

print("Final Generator Loss:", final_gen_loss)
print("Final Discriminator Loss:", final_disc_loss)
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">Generator Test: tensor(-0.4425, grad_fn=&lt;MeanBackward0&gt;)
Discriminator Test: tensor(-0.9662, grad_fn=&lt;MeanBackward0&gt;) tensor(0.3021, grad_fn=&lt;MeanBackward0&gt;)
Generator Loss: tensor(-0.3021, grad_fn=&lt;MulBackward0&gt;)
Discriminator Loss: tensor(1.2683, grad_fn=&lt;SubBackward0&gt;)
Generator Gradient: (tensor([[-0.0995]], grad_fn=&lt;TBackward&gt;), tensor([0.6430], grad_fn=&lt;ViewBackward&gt;))
Discriminator Gradient: (tensor([[[-1.6031]]], grad_fn=&lt;AddBackward0&gt;), tensor([[-1.3847]], grad_fn=&lt;AddBackward0&gt;))
Final Generator Loss: tensor(0.3486, grad_fn=&lt;AddBackward0&gt;)
Final Discriminator Loss: tensor(3.3866, grad_fn=&lt;AddBackward0&gt;)
</code></pre>

username: import-antigravity
cooked: <p>Also just for reference <a href="https://arxiv.org/abs/1802.05642" rel="noopener nofollow ugc">this</a> is the paper I’m trying to implement</p>

username: import-antigravity
cooked: <p>Ah, I figured it out! The problem is that I’m using the pytorch-lightning package, and it was freezing the weights of the other model for each step. On the discriminator step, the gradient norm is actually fully a function of the <em>generator</em> weights, so that was causing the problem. So it looks like I’ll have to probably write a custom optimizer to get the functionality I want.</p>