Default weight initialisation for Conv layers (including SELU)

aroibu1: <p>Firstly, apologise if these are silly questions!</p>
<p>1 I am wondering what is the default initialisation utilised for Conv layers and is this dependent on the nonlinearity selected for after the layer?<br>
2. When using a SELU nonlinearity, does the network automatically initialise the weights using the LeCun Normal Initialisation? If not, how could I implement weight initialisation manually to use the LeCun Normal?</p>
<p>A bit of context:</p>
<p>Reading through the various blog posts and questions from the past few years, for (1) I managed to find two opposing opinions: either that PyTorch automatically initialises all weights to LeCun Normal, or that PyTorch initialises weights based on the non-linearity used after the Conv Layer (Xavier for Tanh and Kaiming He for ReLU and ReLU derivated). However, when I check the source code (<a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/conv.py" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/conv.py</a>), it appears that the default weight initlisation is Kaiming:</p>
<pre><code class="lang-auto">    def reset_parameters(self) -&gt; None:
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)
</code></pre>
<p>In this case, my understanding is that that Kaiming is the default weight initialisation for Conv layers, no matter the following nonlinearity? Thus, most of the posts I have previously read are outdated.</p>
<p>Then, in order to implement the LeCun Normal initialisation, do I need to rewrite reset_parameters in my own code so that it overwrites the default PyTorch code?</p>

Nikronic: <p>Hi,</p>
<p>For the first question, please see these posts:</p>
<ol>
<li><a href="https://discuss.pytorch.org/t/clarity-on-default-initialization-in-pytorch/84696" class="inline-onebox">Clarity on default initialization in pytorch</a></li>
<li><a href="https://discuss.pytorch.org/t/cnn-default-initialization-understanding/85353/2" class="inline-onebox">CNN default initialization understanding</a></li>
</ol>
<p>I have explained the magic number <code>math.sqrt(5)</code> so you can also get the idea behind the relation between non-linearity and init method. Acuatlly, default initialization is uniform.</p>
<p>Also, see this reply in the github thread about it <a href="https://github.com/pytorch/pytorch/issues/15314#issuecomment-477448573" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/15314#issuecomment-477448573</a></p>
<p>About the second question, you can reinitialize weights after initializing them using default values. To do so, you can create your init function similar to available cases in <code>torch.nn.init</code> package and use a code similar to following snippet:</p>
<pre><code class="lang-auto">def init_weights(m):
    """
    Initialize weights of layers using Kaiming Normal (He et al.) as argument of "Apply" function of
    "nn.Module"
    :param m: Layer to initialize
    :return: None
    """

    if isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out')
        nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.BatchNorm2d):
        nn.init.constant_(m.weight, 1)
        nn.init.constant_(m.bias, 0)

model.apply(init_weights)
</code></pre>
<p>Bests</p>

aroibu1: <p>Thanks <a class="mention" href="/u/nikronic">@Nikronic</a>! This makes things slightly more clear to me. However, I still have some questions, if that’s okay?</p>
<p>I understand the use of <strong>math.sqrt(5)</strong> and how this ties in with the nonlinearity, IF the nonlinearity is ReLU or LeakyReLU. In my case, I am using a PReLU non-linearity for now. This is similar to LeakyReLU, and the Kaiming Loss has been created for it. Given this, would the <strong>math.sqrt(5)</strong> still be a good choice for the a-parameter in Kaiming Loss?</p>
<p>Also, thank you for the advise on reinitilisation!</p>

Nikronic: <p>Great, If you are using anything else rather than LeakyReLU, you need to get proper value based on <code>gain</code> of that particular activation function.</p>

aroibu1: <p>I understand. Do you have any suggestions on how this can be done, or any sources that I can consult? I had a look at <em>torch.nn.init.calculate_gain</em> but ‘prelu’ is not a supported nonlinearity.</p>

Nikronic: <p>Yes, you need to find gain yourself in this case. Actually, I am not familiar with calculating gain so I cannot help with that. Let me know if you found anything.</p>

aroibu1: <p>Okay - this is a continuation of the message stream I had with <a class="mention" href="/u/nikronic">@Nikronic</a>, and is my solution to calculating the gain in order to properly use PReLU nonlinearity. I have not yet implemented this myself, but it’s what makes sense to me after reading the derivation in the original paper:  <a href="https://arxiv.org/pdf/1502.01852.pdf" rel="noopener nofollow ugc">https://arxiv.org/pdf/1502.01852.pdf</a>. I’m writing this both for the community and my own later use <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>So, we start from the std definition in the paper:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/2/c2b62d5f7589b16dddf81e3d4ef9bfffe8b77d85.png" alt="Screenshot 2020-06-26 at 15.46.44" data-base62-sha1="rMuWpOqh7iSKCXXBgT8DNYKx2ux" width="273" height="62"></p>
<p>If we rearrange this, we obtain that the standard deviation (=sqrt(Var)) is given by the following two, which are the values obtained from <a href="https://pytorch.org/docs/stable/nn.init.html" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.init — PyTorch 2.1 documentation</a>. The gain is basically identically calculated to the one in LeakyReLU.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/f/0ffb7e1904ac988fa8af03a66f6a66cb3dae0352.png" alt="Screenshot 2020-06-26 at 15.51.30" data-base62-sha1="2hnYUlKfyD2dl6gsN7YUG6T1KH8" width="182" height="60"></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/7/c7737fc70c4f6afe649eb8b49c5cb12309179c13.png" data-download-href="https://discuss.pytorch.org/uploads/default/c7737fc70c4f6afe649eb8b49c5cb12309179c13" title="Screenshot 2020-06-26 at 15.51.54"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c7737fc70c4f6afe649eb8b49c5cb12309179c13_2_377x500.png" alt="Screenshot 2020-06-26 at 15.51.54" data-base62-sha1="ssqthM6fLJ8VVbNf4X3q74ptL5V" width="377" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c7737fc70c4f6afe649eb8b49c5cb12309179c13_2_377x500.png, https://discuss.pytorch.org/uploads/default/original/3X/c/7/c7737fc70c4f6afe649eb8b49c5cb12309179c13.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/c/7/c7737fc70c4f6afe649eb8b49c5cb12309179c13.png 2x" data-dominant-color="F7F7F8"><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 2020-06-26 at 15.51.54</span><span class="informations">389×515 12 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>The equation derived in the original paper is the one used for the normal distribution, not the uniform one  <code>torch.nn.init.kaiming_normal_</code> . So, in theory, we could just use this one. However, this is not possible, as the <code>kaiming_normal_</code> function in PyTorch calls <code>torch.nn.init.calculate_gain</code> which does not accept PReLU as a nonlinearity. Thus, we need a workaround this issue:</p>
<p>The alternative is to just calculate our own standard deviation, which is actually easier than I thought. In the paper, they suggest initiating the <strong>negative_slope</strong> to whatever value we also use to initiate it in our PReLU. For PyTorch, that would be <em><strong>0.25</strong></em> (<a href="https://pytorch.org/docs/stable/nn.html#prelu" class="inline-onebox" rel="noopener nofollow ugc">torch.nn — PyTorch 2.1 documentation</a>). We also need to calculate the <strong>fan_mode</strong>, for which we can look at how this is calculated in the PyTorch source (<a href="https://pytorch.org/docs/stable/_modules/torch/nn/init.html#kaiming_normal_" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/<em>modules/torch/nn/init.html#kaiming_normal</em></a>):</p>
<pre><code class="lang-auto">def _calculate_fan_in_and_fan_out(tensor):
    dimensions = tensor.dim()
    if dimensions &lt; 2:
        raise ValueError("Fan in and fan out can not be computed for tensor with fewer than 2 dimensions")

    num_input_fmaps = tensor.size(1)
    num_output_fmaps = tensor.size(0)
    receptive_field_size = 1
    if tensor.dim() &gt; 2:
        receptive_field_size = tensor[0][0].numel()
    fan_in = num_input_fmaps * receptive_field_size
    fan_out = num_output_fmaps * receptive_field_size

    return fan_in, fan_out
</code></pre>
<p>Following the calcualtion of the std, and knowing we have a mean of 0, we can rewrite the code provided here (Example 10 - <a href="https://www.programcreek.com/python/example/107693/torch.nn.PReLU" class="inline-onebox" rel="noopener nofollow ugc">Python Examples of torch.nn.PReLU</a>) to produce our required weight initialisation:</p>
<pre><code class="lang-auto">        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = fan_in
                negative_slope = 0.25
                m.weight.data.normal_(0, math.sqrt(2. / (n * (1 + negative_slope ** 2)) ) )
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_() 
</code></pre>
<p>Of course, if one uses to use fan_out, they can just replace it using the code above.</p>
<p><a class="mention" href="/u/nikronic">@Nikronic</a> - does what I wrote above make sense? Asking as I think you have more experience with these things than I do.</p>

Nikronic: <p>Hi, Sorry for my late answer, I am struggling with final exams! <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>
<p>First, thank you for your deep explanation. Secondly, I do not have strong mathematical background, so I do not think I am eligble to validate this but sounds fine to me.</p>
<p>Another point I would like to mention is that PyTorch uses uniform for initializing weights in convs and linear layers so if gain in PReLU is identical to LeakyReLU, then to achieve the range of <code>[-1/sqrt(fan_mode), 1/sqrt(fan_mode)]</code> for uniform distribution, still we need to consider <code>negative_slope=sqrt(5)</code> where otherwise it will lead to a different scenario.</p>
<p>I think we need to discuss this as a feature request so the main developers can help us with that. So, I think it would be great idea to create an issue on github.<br>
Here is another <a href="https://github.com/pytorch/pytorch/issues/24991" rel="noopener nofollow ugc">issue</a> related to this idea that may help. Furthermore, this thread considered another perspective which I have no clue what they are talking about <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20">.</p><aside class="quote" data-post="1" data-topic="20854">
  <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/jpeg729/48/3464_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/calculate-gain-tanh/20854">Calculate_gain('tanh')</a> 
  </div>
  <blockquote>
    I understand the results returned by calculate_gain for linear, relu, leaky_relu and sigmoid. 
Can anyone tell me why calculate_gain('tanh') returns 5/3 ?
  </blockquote>
</aside>

<p>If you created the issue on github, could you please also tag me so I can keep track of things? my username is <a href="https://github.com/Nikronic" rel="noopener nofollow ugc">Nikronic</a></p>
<p>Thank you</p>