Converting from nn.BatchNorm2d to nn.LayerNorm in CNN

asberman: <p>For improved Wasserstein GAN (aka Wasserstein GAN with gradient penalty [WGAN-GP]), layer normalization is recommended in the discriminator, as opposed to <code>nn.BatchNorm2d</code></p>
<p>I see that <code>nn.LayerNorm</code> was (relatively) recently added to <code>torch.nn.modules</code>, and I’d like to use it, as opposed to writing my own layer normalization. However, just replacing calls to <code>nn.BatchNorm2d(input_size)</code>  with <code>nn.LayerNorm(input_size)</code> gives the following error:</p>
<pre><code>File "/usr/local/lib/python3.6/site-packages/torch/nn/functional.py", line 1314, in layer_norm
PSTtorch.backends.cudnn.enabled)
PSTRuntimeError: Given normalized_shape=[128], expected input with shape [*, 128], 
but got input of size[128, 128, 16, 16]
</code></pre>
<p>This is the structure of the network:</p>
<pre><code>class WgangpDiscriminator(nn.Module):  
    def __init__(self, channels, num_disc_filters):
        super(WgangpDiscriminator, self).__init__()
        self.ngpu = 1
        self.main = nn.Sequential(
        #num_disc_filters = 64
        nn.Conv2d(channels, num_disc_filters, 4, 2, 1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),

        nn.Conv2d(num_disc_filters, num_disc_filters * 2, 4, 2, 1, bias=False),
        nn.LayerNorm(num_disc_filters * 2),
        nn.LeakyReLU(0.2, inplace=True),

        nn.Conv2d(num_disc_filters * 2, num_disc_filters * 4, 4, 2, 1, bias=False),
        nn.LayerNorm(num_disc_filters * 4),
        nn.LeakyReLU(0.2, inplace=True),

        # state size. (num_disc_filters*4) x 8 x 8
        # nn.Conv2d(num_disc_filters * 4, num_disc_filters * 8, 4, 2, 1, bias=False),
        # nn.BatchNorm2d(num_disc_filters * 8),
        # nn.LeakyReLU(0.2, inplace=True),

        # was num_disc_filters * 16
        nn.Conv2d(num_disc_filters * 4, num_disc_filters * 8, 4, 2, 1, bias=False),
        nn.LayerNorm(num_disc_filters * 8),
        nn.LeakyReLU(0.2, inplace=True),
        
        # was num_disc_filters * 16
        nn.Conv2d(num_disc_filters * 8, 1, 4, 1, 0, bias=False),
        # ty  https://github.com/pytorch/examples/issues/70 apaske
    )

    def forward(self, inp):
         if isinstance(inp.data, torch.cuda.FloatTensor) and self.ngpu &gt; 1:
             output = nn.parallel.data_parallel(self.main, inp, range(self.ngpu))
         else:
             output = self.main(inp)
         return output.view(-1, 1).squeeze(1)
</code></pre>
<p><code>nn.LayerNorm</code> expects <code>normalized_shape</code> as input ( an <code>int</code>, <code>list</code> or <code>torch.Size</code>), but <code>nn.Conv2d</code> layers don’t have <code>.size</code> , <code>.get_shape()</code> or <code>.shape()</code>, so I can’t follow the example in the docs:</p>
<pre><code>input = torch.randn(20, 5, 10, 10) 
# With Learnable Parameters 
m = nn.LayerNorm(input.size()[1:])
</code></pre>
<p>How do I do this conversion?</p>
<p>Many thanks in advance!</p>

b037059c4e221d6caee8: <p>Hi,I meet the same problem, do you solve this problem? THX</p>

11157: <p>I also meet this problem, do you solve this? Thanks.</p>

AboMAd: <p>Hi asberman,</p>
<p>As I understand LayerNorm will compute mean and variance elementwise (not per batch), thus you should pass the spatial dimension of the input, not the channel dimension as in the case of BatchNorm.<br>
Actually, I am doing the same work, and you can try to change the following:<br>
the first layer norm :<br>
nn.LayerNorm(num_disc_filters * 2), --&gt; nn.LayerNorm([num_disc_filters * 2, 16, 16]),</p>
<p>the second:<br>
nn.LayerNorm(num_disc_filters * 4), --&gt; nn.LayerNorm([num_disc_filters * 4, 8, 8]),</p>
<p>the third:<br>
nn.LayerNorm(num_disc_filters * 8), --&gt; nn.LayerNorm([num_disc_filters * 8, 4, 4]),</p>
<p>I think this will work without error.<br>
However, I found that even after replacing the BatchNorm with the LayerNorm, the model still diverges!!!<br>
If you have solved this properly, please inform me.</p>
<p>Thanks</p>

sharkdeng: <p>I found your answer is like hard coding, how to use an unify way?</p>