How to pack two modules of GAN(generator and discriminator) in a single module?

username: Jeet
cooked: <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.LeakyReLU(0.01),

            nn.Linear(128, 1),

            nn.Sigmoid(),

        )

    def forward(self, x):

        return self.disc(x)

class Generator(nn.Module):

    def __init__(self, z_dim, img_dim):

        super().__init__()

        self.gen = nn.Sequential(

            nn.Linear(z_dim, 256),

            nn.LeakyReLU(0.01),

            nn.Linear(256, img_dim),

            nn.Tanh(), 

        )

    def forward(self, x):

        return self.gen(x) 
</code></pre>
<p>How should I design the forward method in that case if the generator and the discriminator both are considered a single module?</p>

username: KFrank
cooked: <p>Hi Jeet!</p>
<aside class="quote no-group" data-username="Jeet" data-post="1" data-topic="147564" 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/jeet/48/47500_2.png" class="avatar"> Jeet:</div>
<blockquote>
<p>How should I design the forward method in that case if the generator and the discriminator both are considered a single module?</p>
</blockquote>
</aside>
<p>First, as I am sure you are aware, there is no problem with having one<br>
<code>Module</code> contain other <code>Module</code>s.</p>
<p>Now what your <code>forward()</code> method should look like depends on the<br>
details of how you want to train your GAN.  (There are a number of<br>
approaches with varying degrees of nuance.)</p>
<p>In one simple scheme, you perform a forward pass in which you both<br>
generate a fake image in your <code>Generator</code> <code>Module</code> and then pass<br>
that fake image (possibly together with a real image) through your<br>
<code>Discriminator</code> <code>Module</code> on to the final loss function that rewards your<br>
<code>Discriminator</code> for successfully distinguishing fake from real.</p>
<p>The issue is that you also want to reward your <code>Generator</code> for successfully<br>
fooling  your <code>Discriminator</code> into mistaking fake for real.  That is, you<br>
want to train – in the same backward pass – your <code>Discriminator</code> and<br>
<code>Generator</code> “in opposite directions.”  One way to do this is to flip the sign<br>
of the gradients that are backpropagated back through the <code>Generator</code>.</p>
<p>You can to this by introducing (after, of course, having written it) a layer,<br>
packaged as a custom <code>torch.autograd.Function</code>, that passes the tensor<br>
through unchanged on the forward pass, but that then flips the sign of the<br>
gradient on the backward pass.</p>
<p>You would then do something like this:</p>
<pre data-code-wrap="python"><code class="lang-python">class MyGANModule (nn.Module) :
    def __init__ (self, in_features, z_dim, ing_dim) :
        super().__init__()
        self.my_gen = Generator (z_dim, img_dim)
        self.gradient_flipper = CustomGradientSignFlippingFunction()
        self.my_disc = Discriminator (in_features)
    
    def forward (self, x) :
        x = self.my_gen (x)
        x = self.gradient_flipper.apply (x)
        x = self.my_disc (x)
        return  x
</code></pre>
<p>(You could, alternatively, package <code>Generator</code>,<br>
<code>CustomGradientSignFlippingFunction</code>, and <code>Discriminator</code> together<br>
as a <code>Sequential</code>, if you prefer the <code>Sequential</code> style.)</p>
<p>So <code>my_gen</code> and <code>my_disc</code> are themselves both <code>Module</code>s, but they are<br>
contained in (said another way, are properties of) your <code>MyGANModule</code><br>
<code>Module</code>.</p>
<p><code>MyGANModule</code> has its own <code>forward()</code> method, but when that <code>forward()</code><br>
method calls, for example, <code>my_gen (x)</code>, pytorch’s infrastructure then<br>
calls <code>Generator</code>’s own <code>forward()</code> method on <code>x</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Jeet
cooked: <p>Hi Frank,</p>
<p>Thanks a lot again for the detailed explanation. You are correct that this flipping of gradient is required if we need to train the Generator in the same backward pass. In this regard let me tell you, that I am using a different formula to train the GAN. The weights of the genrators and the Discriminator will be updated taking help of that gradient descent formula. However, the formula requires an input which is a Jacobian matrix. The Jacobian matrix will contain the first order partial derivative of the loss functions of the generator and the discriminator with respect to all the weights (Both Generator and Discriminator).<br>
Hence, I was thinking that if I can have both the modules of generator and the  discriminator in a single module then I could have the partial derivative of the respective loss functions with respect to all the weights. For example- LossGenerator.backward() and LossDiscriminator.backward() but with repect to all the weights in the model.<br>
However, the problem here is in GAN, the losses are calculated seperately for the generator and the discriminator. It is because the real data(real image or img_dim above) is fed only to the discriminator but not to the generator. The output of the discriminator is fed to the generator and the loss of the generator is calculated. Since, the loses are calculated seperately, I have no clue how can I calculate the first order partial derivative for each of the loss functions with respect to all the weights.<br>
In the MyGANModule module suggested by you above -<br>
should I have 3 forward method?One for the generator, one for the discriminator and the other one combined (as you have written above)? But I am not sure if that is going to help. If I use MyGANModule to train the discriminator with the real image and find the loss, then can I find the gradients(1st order partial derivative) with respect to the weights of the generator as well?<br>
Could you please provide a code snippet as an example for the same?</p>