Adding new layers between intermediate layers of a GAN

avalon1511: <p>I have a Generator defined as:</p>
<pre><code class="lang-auto">class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        return self.main(input)
</code></pre>
<p>Between every <code>ConvTranspose2d</code> layer and <code>BatchNorm2d</code> layer I want to insert some kind of mathematical operation that is applied on the output of the conv layer.</p>
<p>For example, let’s say the output of the conv layer is <code>X</code> and i want to apply a shift <code>b</code> to this output such that <code>X + b </code> is now the input to the BN layer.</p>
<p>How can I do this?</p>

ptrblck: <p>You could create a new custom <code>nn.Module</code> and add it into the <code>nn.Sequential</code> container:</p>
<pre><code class="lang-python">class ShiftLayer(nn.Module):
    def __init__(self, b):
        super().__init__()
        self.register_buffer('b', b)
        
    def forward(self, x):
        return x + self.b

class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
            ShiftLayer(b=torch.tensor([2.])),
            nn.BatchNorm2d(ngf * 8),
            ...
</code></pre>
<p>Registering <code>b</code> as a buffer will make sure that the <code>to()</code> operation such as <code>model.to('cuda')</code> will move this buffer to the specified device etc.<br>
Also, it will be stored in the <code>state_dict</code> of the module so that you could save and load it properly.</p>

avalon1511: <p>Hi, that solution works, thanks! I’m now trying to modify the x and b parameters from this ShiftLayer based on some conditions</p>
<p><code>netG.x</code> and <code>netG.b</code> does not work, it gives me the error</p>
<pre><code class="lang-auto">torch.nn.modules.module.ModuleAttributeError: 'Generator' object has no attribute 'x'

</code></pre>
<p>So I tried accessing them through <code>netG.module.x</code> but that does not work either and I get the error:</p>
<pre><code class="lang-auto">torch.nn.modules.module.ModuleAttributeError: 'Generator' object has no attribute 'module'

</code></pre>
<p>How can I do this?</p>

ptrblck: <p><code>ShiftLayer</code> doesn’t have an <code>x</code> parameter, but you could access the <code>b</code> buffer via:</p>
<pre><code class="lang-python">netG.main[1].b
</code></pre>
<p>assuming you are using the model definition posted in my code snippet.</p>

avalon1511: <p>Yes, sorry that was my bad, I am wanting to update <code>b</code>. So I actually tried what you suggested:</p>
<pre><code class="lang-auto">netG.main[1].b  += 5
</code></pre>
<p>and now the error is</p>
<pre><code class="lang-auto">'ShiftLayer' object has no attribute 'b'

</code></pre>
<p>Although I have initialized the generator exactly as you posted in your code snippet.</p>

avalon1511: <p>Okay, I looked over my code and it was a naming error on my end. It’s fixed now! thanks! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

avalon1511: <p>Is it possible to define <code>ShiftLayer</code> so that I can train its parameter b with a <code>torch.optim()</code> function?</p>

ptrblck: <p>Yes, register <code>b</code> as an <code>nn.Parameter</code> and it’ll be trainable:</p>
<pre><code class="lang-python">self.b = nn.Parameter(torch.tensor(...))
</code></pre>

avalon1511: <p>This has to be done when I define the <code>ShiftLayer</code> right?</p>
<pre><code class="lang-auto">class ShiftLayer(nn.Module):
    def __init__(self, b):
        super().__init__()
        self.register_buffer('b', b)
        self.b = nn.Parameter(torch.tensor(...))
        
    def forward(self, x):
        return x + self.b
</code></pre>
<p>like this?</p>

ptrblck: <p>Yes, you can define the parameter in the <code>ShiftLayer</code>, but would need to remove the <code>self.register_buffer</code> as it’s not needed.</p>

avalon1511: <p>So my final definition for <code>ShiftLayer</code> is</p>
<pre><code class="lang-auto">class ShiftLayer(nn.Module):
    def __init__(self, b):
        super().__init__()
        self.b = nn.Parameter(torch.tensor(...))
        
    def forward(self, x):
        return x + self.b
</code></pre>
<p>However, this gives me an error:</p>
<pre><code class="lang-auto"> self.b = nn.Parameter(torch.tensor(...)),                                                                                                ^                                                                                                                               
SyntaxError: expression cannot contain assignment, perhaps you meant "=="?  
</code></pre>
<p>Not sure how else I would assign <code>b</code> as a parameter</p>

ptrblck: <p>The <code>...</code> was just a placeholder for your actual values, so you would need to fill <code>torch.tensor</code> either with your values or use a factory method given a shape, e.g. <code>torch.randn(10, 10)</code>.</p>

avalon1511: <p>yeah  of couse. I have actual values like so -</p>
<pre><code class="lang-auto">self.b = nn.Parameter(torch.randn(32)),
</code></pre>
<p>still getting this error though <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ptrblck: <p>I cannot reproduce the issue using:</p>
<pre><code class="lang-python">class ShiftLayer(nn.Module):
    def __init__(self, b):
        super().__init__()
        self.b = nn.Parameter(torch.tensor(1.))
        
    def forward(self, x):
        return x + self.b
    
layer = ShiftLayer(1)
out = layer(torch.randn(1))
</code></pre>
<p>so could you post an executable code snippet raising this issue?</p>

avalon1511: <p>Sorry for the late reply, I fixed it! I had an accidental <code>Sequential()</code> command right after <code>def __init__()</code> which was the problem</p>