Weights are not updated when using custom function before loss calculation

truss712: <p>Hello community,</p>
<p>I am trying to implement a model-based GAN, where I try to predict simulation parameters from the generator and use the generator’s output to calculate simulations from a custom monte Carlo simulation function. I provide these simulated images to the discriminator with true images to calculate the loss.</p>
<p>I then try to update both generator and discriminator based on the losses calculated in the above stage.</p>
<p>The problem with this is, weights of the generator are not updating.</p>
<ol>
<li>
<p>Is this approach feasible for PyTorch, where I have to break the graph on purpose?</p>
</li>
<li>
<p>Is there a better way to address this issue?</p>
</li>
<li>
<p>Is it possible if I define my simulation script as a custom torch layer?</p>
</li>
</ol>
<p>I am attaching a simple example flow for reference.</p>
<pre><code class="lang-auto">class Discriminator(nn.Module):
    def __init__(self, in_features):
        super().__init__()
        self.disc = nn.Sequential(
            nn.Linear(in_features, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )

    def forward(self, x):
        return self.disc(x)


class Generator(nn.Module):
    def __init__(self, n_params, scalar, test_transform=None):
        super().__init__()
        self.scalar = scalar
        self.test_transform = test_transform

        self.gen = nn.Sequential(
            nn.Linear(n_params, 2 * n_params),
            nn.ReLU(),
            nn.Linear(2 * n_params, 2 * n_params),
            nn.ReLU(),
            nn.Linear(2 * n_params, n_params),
            nn.Tanh(),
        )

    def forward(self, x):
        p1 = self.gen(x)
        params = self.scalar.inverse_transform(p1.clone().detach().numpy())
        params = p1 + torch.tensor([100, 1.0e-12, 100, 1e-3, -1 * math.pi / 180])
        # Generate simulation
        im = simulate_image(params[0])
        fake = self.test_transform(image=np.swapaxes(np.array(im), 1, -1))
        fake = torch.from_numpy(fake["image"])
        fake = fake.contiguous()
        return fake
</code></pre>
<p>I also tried without converting my p1 Tensor to numpy and directly giving tensor to my simulation function. But that also didn’t work.</p>

KFrank: <p>Hi Truss!</p>
<aside class="quote no-group" data-username="truss712" data-post="1" data-topic="123259" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/truss712/48/39068_2.png" class="avatar"> truss712:</div>
<blockquote>
<ol>
<li>Is this approach feasible for PyTorch, where I have to break the graph on purpose?</li>
</ol>
</blockquote>
</aside>
<p>Probably not.  To make pytorch’s autograd and gradient-descent-based<br>
optimizers work, you would have to “unbreak” the graph, patching it<br>
back together somehow.</p>
<aside class="quote no-group">
<blockquote>
<ol start="2">
<li>Is there a better way to address this issue?</li>
</ol>
</blockquote>
</aside>
<p>The core problem is that it is hard to make the result of a Monte Carlo<br>
simulation differentiable with respect to your model parameters.</p>
<p>Consider the derivative of the area of a circle with respect to its radius.<br>
Now consider calculating the area with a Monte Carlo “simulation”<br>
(integration).  Sample two-dimensional points uniformly from the square<br>
<code>x in [-1.0, 1.0]; y in [-1.0, 1.0]</code> and count the fraction of<br>
those points that fall in a circle of radius <code>r &lt;= 1.0</code>.  You then have<br>
<code>area = 4 * fraction</code>.  In this scheme, how would you calculate the<br>
derivative of <code>area</code> with respect to <code>r</code>?</p>
<p>You can imagine generating “simple” random numbers (say uniform<br>
or normal), using them to drive a Monte-Carlo process that is<br>
implemented with pytorch differentiable tensor functions, and then<br>
using autograd to compute the gradient of some final loss function<br>
with respect to parameters that characterize the Monte-Carlo process<br>
(as well as with respect to the random numbers, if you so choose).</p>
<p>But whether you can do this in practice depends on the details of<br>
your Monte-Carlo process.</p>
<aside class="quote no-group">
<blockquote>
<ol start="3">
<li>Is it possible if I define my simulation script as a custom torch layer?</li>
</ol>
</blockquote>
</aside>
<p>Yes, in principle.  (But it may be impractical or impossible for your use<br>
case.)</p>
<p>You would do this by defining a <a href="https://pytorch.org/docs/1.8.1/autograd.html#torch.autograd.Function" rel="noopener nofollow ugc">custom autograd function</a> (your “custom<br>
torch layer”).</p>
<p>It will have a <code>forward()</code> method that can do most anything – numpy<br>
calculations, a Monte Carlo simulation, any sort of black box.  But it<br>
also has to have a matching <code>backward()</code> method that calculates the<br>
gradient (or an approximation to, or surrogate for the gradient) of the<br>
output of <code>forward()</code> with respect to its input.</p>
<p>It’s this <code>backward()</code> method that “patches the graph back together.”</p>
<p>Best.</p>
<p>K. Frank</p>

truss712: <p>Thank you <a class="mention" href="/u/kfrank">@KFrank</a> for the detailed answer.<br>
My assumptions are also similar for 1 and 2.<br>
I will try approach 3 and it might work as you suggested.<br>
I don’t require gradients for the simulation layer as I just want to update params prediction network.<br>
Does it make sense if I give precalculated dummy gradients in backward function of the custom layer?<br>
Thank you in advance.!</p>

KFrank: <p>Hi Truss!</p>
<aside class="quote no-group" data-username="truss712" data-post="3" data-topic="123259" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/truss712/48/39068_2.png" class="avatar"> truss712:</div>
<blockquote>
<p>I don’t require gradients for the simulation layer as I just want to update params prediction network.</p>
</blockquote>
</aside>
<p>I don’t really understand what you are saying here.</p>
<p>But let me ask this question:  Should the update you apply to your<br>
prediction-network parameters depend <em>in any way</em> on the output<br>
of the simulation layer?  If not, I can’t see any reason that you would<br>
need to run the simulation layer to perform the prediction-network<br>
update.  If so, it would seem that you would have to backpropagate<br>
gradients of your final loss function (that depends on your simulation<br>
layer) through the simulation layer back to the prediction network.</p>
<aside class="quote no-group">
<blockquote>
<p>Does it make sense if I give precalculated dummy gradients in backward function of the custom layer?</p>
</blockquote>
</aside>
<p>This doesn’t make sense to me.  If your custom simulation layer<br>
matters, dummy gradients will obscure its structure and make it not<br>
matter to the prediction-layer update.  (And if your simulation layer<br>
doesn’t matter, you should be able to update your prediction layer<br>
based solely on its output without passing anything – forward or<br>
backward – through the simulation layer.)</p>
<p>At a conceptual level – ignoring gradients and specific optimization<br>
algorithms – your loss function is supposed to tell you how good a<br>
job the parameters of your prediction network are doing.  What role<br>
does your simulation play in telling you whether your prediction<br>
network is performing well or poorly?  Or to ask the same question<br>
more concretely, how does your final loss function depend on your<br>
simulation layer and its output?</p>
<p>Best.</p>
<p>K. Frank</p>

truss712: <aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="123259">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Should the update you apply to your<br>
prediction-network parameters depend <em>in any way</em> on the output<br>
of the simulation layer?</p>
</blockquote>
</aside>
<p>The flow of my network is in a way where parameters are learnt in neural network. In a simple way, it can be compared with normal GAN where generator here does not generate image but just bunch of parameters. These parameters are updated by calculating losses by comparing simulated image with true image in discriminator.</p>
<aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="123259">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>What role<br>
does your simulation play in telling you whether your prediction<br>
network is performing well or poorly?</p>
</blockquote>
</aside>
<p>Simulator generates images from learnt parameters from my prediction network (Generator), and this simulated image from my simulation algorithm is then compared with true image in discriminator. Losses are also calculated here. So, in a way, network is not important at all in simulation stage. I hope this answers it. But when I tried developing network like in my sample code, it is not updating any weights for my prediction parameters part. (Generator class weights)</p>
<p>To summarize,</p>
<pre><code class="lang-auto">gen = Generator()
learnt_params = gen(noisy_input)
#Generate simulated image
fake= simulate_image(learnt_params)
disc_real = disc(real).view(-1)
lossD_real = criterion(disc_real, torch.ones_like(disc_real))
disc_fake = disc(fake).view(-1)

 lossD_fake = criterion(disc_fake, torch.zeros_like(disc_fake))
 lossD = (lossD_real + lossD_fake) / 2
 disc.zero_grad()
 lossD.backward(retain_graph=True)
 opt_disc.step()
### Train Generator: min log(1 - D(G(z))) &lt;-&gt; max log(D(G(z))
output = disc(fake).view(-1)
lossG = criterion(output, torch.ones_like(output))

</code></pre>
<p>I hope the issue is more clear now.<br>
Thank yo</p>
<aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="123259">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Should the update you apply to your<br>
prediction-network parameters depend <em>in any way</em> on the output<br>
of the simulation layer?</p>
</blockquote>
</aside>
<p>Yes, simulated image generated using simulation layer (only numpy functions) will be compared with real image in discriminator stage to train it and then use discriminator to calculate losses for generator.<br>
In simple words, I am using actual GAN structure, but instead of generator generating whole image from noise, I use generator to generate parameters which I use in simulation layer.<br>
I hope this is more clear now.</p>
<aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="123259">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>What role<br>
does your simulation play in telling you whether your prediction<br>
network is performing well or poorly?</p>
</blockquote>
</aside>
<p>Role of simulation is just to simulate image from predicted parameters, but it helps calculating losses as I explained before. I dont have any real known values for parameters and I can only compare real images and simulated images for their distribution matching.<br>
Here is the sample code…</p>
<pre><code class="lang-auto">fake_params = gen(params)
simulated_image = simulate_image(fake_params) #consists of only numpy functions
disc_real = disc(real)
disc_fake = disc(fake)
lossD_real = criterion(disc_real, torch.ones_like(disc_real))
lossD_fake = criterion(disc_fake, torch.zeros_like(disc_fake))
lossD = (lossD_real + lossD_fake) / 2
disc.zero_grad()
lossD.backward(retain_graph=True)
opt_disc.step()
output = disc(fake)
lossG = criterion(output, torch.ones_like(output))
gen.zero_grad()
lossG.backward()
opt_gen.step()
</code></pre>
<p>I hope now the connection and role of the simulator are more clear.</p>

KFrank: <p>Hi Truss!</p>
<aside class="quote no-group quote-modified" data-username="truss712" data-post="5" data-topic="123259" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/truss712/48/39068_2.png" class="avatar"> truss712:</div>
<blockquote>
<aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="123259">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Should the update you apply to your<br>
prediction-network parameters depend <em>in any way</em> on the output<br>
of the simulation layer?</p>
</blockquote>
</aside>
<p>…<br>
These parameters are updated by calculating losses by comparing simulated image with true image in discriminator.<br>
…<br>
this simulated image from my simulation algorithm is then compared with true image in discriminator. Losses are also calculated here.<br>
…<br>
Yes, simulated image generated using simulation layer (only numpy functions) will be compared with real image in discriminator stage to train it and then use discriminator to calculate losses for generator.</p>
</blockquote>
</aside>
<p>Reading between the lines, you are, as I understand it, trying to train<br>
the discriminator to tell simulated images apart from real ones, and to<br>
train the generator to generate simulation parameters that lead to<br>
simulated images that are hard to tell apart from real ones.</p>
<p>So, yes, in order to update the generator “weights” you do need to<br>
backpropagate the final discriminator loss back through the simulator.<br>
This is the only way the generator knows whether it did a good job<br>
of generating parameters that produce simulated images that look<br>
realistic to the discriminator.</p>
<p>The simulation process is an important, non-trivial part of how the<br>
generator “weights” produce the simulated images that are used to<br>
calculate the final loss, so to train the generator “weights” you have<br>
to backpropagate through it properly.</p>
<p>Best.</p>
<p>K. Frank</p>