How to reshape the output of a Linear layer in Pytorch?

username: Chris_Davidson
cooked: <p>Hi, I’m implementing Generator of a GAN and I need to reshape output of Linear Layer to particular dimension, here is the image of the Generator.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/6/3672990ad1ea7b31d5ee0713e66b838901e3091f.png" data-download-href="https://discuss.pytorch.org/uploads/default/3672990ad1ea7b31d5ee0713e66b838901e3091f" title="Screenshot 2022-10-21 102802"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/6/3672990ad1ea7b31d5ee0713e66b838901e3091f_2_690x463.png" alt="Screenshot 2022-10-21 102802" data-base62-sha1="7LFj6lQldHLD9r20V1Oqg781fX1" width="690" height="463" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/6/3672990ad1ea7b31d5ee0713e66b838901e3091f_2_690x463.png, https://discuss.pytorch.org/uploads/default/original/3X/3/6/3672990ad1ea7b31d5ee0713e66b838901e3091f.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/6/3672990ad1ea7b31d5ee0713e66b838901e3091f.png 2x" data-dominant-color="F7F7F7"><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">Screenshot 2022-10-21 102802</span><span class="informations">763×512 90.2 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>Here is my code of Generator:</p>
<pre><code class="lang-auto">class Generator(nn.Module):

    def __init__(self, z_dim=100, im_chan=1, hidden_dim=64, rdim=9216):
       
        super(Generator, self).__init__()
        self.z_dim = z_dim
        self.gen = nn.Sequential(
             nn.Linear(z_dim, rdim),
             nn.BatchNorm2d(rdim,momentum=0.9),
             nn.ReLU(inplace=True),
 ----&gt;       torch.reshape(rdim, (6,6,256)), 
             self.make_gen_block(rdim, hidden_dim*2),
             self.make_gen_block(hidden_dim*2,hidden_dim),
             self.make_gen_block(hidden_dim,im_chan,final_layer=True),
        )
   def make_gen_block(self, input_channels, output_channels, kernel_size=1, stride=2, final_layer=False):

        if not final_layer:
           return nn.Sequential(
              nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
              nn.BatchNorm2d(output_channels),
              nn.ReLU(inplace=True)
        )
        else:
              return nn.Sequential(
              nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
              nn.Tanh()
        )

   def unsqueeze_noise(self, noise):
       return noise.view(len(noise), self.zdim, 1, 1)

   def forward(self, noise):
       x = self.unsqueeze_noise(noise)
       return self.gen(x)
</code></pre>
<p>** <strong>The problem is when I try to reshape the output of Linear layer after BatchNorm and ReLU (in fig. Dense as they have used Tensorflow) it is throwing error as :TypeError: reshape(): argument ‘input’ (position 1) must be Tensor, not int</strong>**<br>
torch.reshape is not looking good to me from start but nn does not have reshape in it so used torch explicitly. Can somebody help me to get through this? Thanks.</p>

username: ptrblck
cooked: <p>You could create a custom <code>nn.Module</code> and use the desired <code>view</code> or <code>reshape</code> call in its <code>forward</code> method.</p>

username: Chris_Davidson
cooked: <p>Can your provide some example? I’m new to Pytorch.</p>

username: ptrblck
cooked: <p>Something like this should work:</p>
<pre><code class="lang-python">class Reshape(nn.Module):
    def __init__(self, shape):
        super().__init__()
        self.shape = shape

    def forward(self, x):
        return torch.reshape(x, self.shape)
</code></pre>

username: hduc-le
cooked: <p>In an easy way, I think you can try the following:</p>
<pre><code class="lang-auto">class Generator(nn.Module):
    def __init__(self, z_dim=100, im_chan=1, hidden_dim=64, rdim=9216):
       
        super(Generator, self).__init__()
        self.z_dim = z_dim
        self.hidden_dim = hidden_dim
        self.rdim = rdim
        self.im_chan = im_chan
        
        self.linear = nn.Linear(z_dim, rdim)
        self.batch_norm = nn.BatchNorm2d(rdim,momentum=0.9)
        self.relu = nn.ReLU()
        
        self.block1 = self.make_gen_block(rdim, hidden_dim*2)
        self.block2 = self.make_gen_block(hidden_dim*2,hidden_dim)
        self.block3 = self.make_gen_block(hidden_dim,im_chan, final_layer=True)
        
    def make_gen_block(self, input_channels, output_channels, kernel_size=1, stride=2, final_layer=False):

        if not final_layer:
           return nn.Sequential(
              nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
              nn.BatchNorm2d(output_channels),
              nn.ReLU(inplace=True)
        )
        else:
            return nn.Sequential(
            nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
            nn.Tanh()
        )
        
    def unsqueeze_noise(self, noise):
        return noise.view(len(noise), self.zdim, 1, 1)

    def forward(self, noise):
        x = self.unsqueeze_noise(noise)
        x = self.linear(x)
        x = self.batch_norm(x)
        x = self.relu(x)
        x = x.view(6,6,256)
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        return x

</code></pre>
<p>Let me know if you still face the problem</p>