Replace loss value using torch.no_grad?

username: YamaKruse
cooked: <p>Hi, I have a Generator network that learns matching relationships (e.g., getting index pairs based on matching scores) between two sets, and then I have a Discriminator network judging if the matching quality is good.</p>
<p>The problem is that the discriminator takes matched data as input, and getting index pairs is a non-differentiable operation and will break the graph, so I cannot use the Discriminator loss to update Generator.</p>
<p>One workaround I am considering is that, I can calculate the mean value of the matching scores, which is still differentiable, replace this value with the Discriminator loss value under <code>with.no_grad()</code>, and then do <code>loss.backward()</code> to update the Generator.</p>
<p>So this is more like setting a fake loss function based on the matching scores and then replacing the loss value with the actual loss value I want to use (calculated by Discriminator). My understanding is that it will use the gradient computation graph from the matching scores to update the Generator.</p>
<p>Would you think this idea will work? Are there any misunderstandings or issues I am not aware of? Any suggestions or comments would be greatly appreciated!</p>

username: YamaKruse
cooked: <p>Hello, can someone provide some comments here? Thanks.</p>

username: ptrblck
cooked: <p>I don’t think your idea sounds valid as manipulating the loss value in a <code>no_grad()</code> context should either raise an error or should not change the gradient calculation at all.</p>

username: YamaKruse
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thanks so much for the reply! Do you think there is a way to modify a loss value without breaking the gradient graph (such as detaching it and changing the value)? I am currently doing something like this:</p>
<pre><code class="lang-auto"># update G
generator.zero_grad()

# get matching scores (differentiable)
scores = generator(data, noise)

# get indices from scores (non-differentiable)
indices = scores.max(1).indices

# permutate the data and calculate loss via discriminator 
d_output = discriminator(data[indices])
d_loss = criterion(d_output, d_label)

# make fake loss using scores.mean() and change its value via detaching
fake_g_loss = scores.mean()
fake_g_loss_detached = fake_g_loss.detach()
fake_g_loss_detached *= 0.
fake_g_loss_detached += d_loss.item()

# Calculate gradients for G
fake_g_loss_detached.backward()
  
# Update G
optimizerG.step()
</code></pre>
<p>I checked <code>fake_g_loss_detached</code> and it seems the value was updated (replaced with <code>d_loss</code> value) and the computation graph was not affected. I also trained the model and it is running. But I am not sure if this really works. Would such operations make sense to you?</p>

username: ptrblck
cooked: <p>I’m currently unsure why your code works at all, since <code>fake_g_loss_detached</code> is indeed detached (as the name indicates) and the <code>backward()</code> call should fail as seen here:</p>
<pre><code class="lang-python">generator = nn.Linear(10, 10)
data = torch.randn(1, 10)

discriminator = nn.Linear(10, 10)
criterion = nn.MSELoss()

# get matching scores (differentiable)
scores = generator(data)

# get indices from scores (non-differentiable)
indices = scores.max(1).indices

# permutate the data and calculate loss via discriminator 
d_output = discriminator(data)
d_loss = criterion(d_output, torch.rand_like(d_output))

# make fake loss using scores.mean() and change its value via detaching
fake_g_loss = scores.mean()
fake_g_loss_detached = fake_g_loss.detach()
fake_g_loss_detached *= 0.
fake_g_loss_detached += d_loss.item()

# Calculate gradients for G
fake_g_loss_detached.backward() # this should fail !!!
# RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>

username: YamaKruse
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Oh sorry I made a mistake. It was actually <code>fake_g_loss.backward()</code>. I tried to simplify the code. The original code has more content. Would <code>fake_g_loss.backward()</code> update the network with the new value or will it still use the old one?</p>

username: ptrblck
cooked: <p>Ah OK, this makes sense.<br>
<code>fake_g_loss.backward()</code> will use its own loss value to calculate the gradients. You could seed the code to make it reproducible and check the gradients after the backward call e.g. via:</p>
<pre><code class="lang-python">print(generator.weight.grad.abs().sum())
</code></pre>
<p>Then you could play around with scaling <code>fake_g_loss_detached</code> and would see that no value changes the computed gradients.</p>
<p>I’m not sure, if this would fit your use case, but in case you want to use a different code path in the backward pass, you could check e.g. <a href="https://discuss.pytorch.org/t/how-to-use-an-optimizer-to-update-non-differentiable-parameters/20668/2">this approach</a> described by <a class="mention" href="/u/tom">@tom</a>.</p>

username: YamaKruse
cooked: <p>Ah thanks so much for the verification! Yes I tried a bit and it looks like it only uses its own value to calculate gradients. Also thanks for sharing with me the approach. It looks promising and I will give it a try.</p>