Initialize the weights in layers -- Which method is most recommended?

Nick_ishere: <p>Dear experienced friends,</p>
<p>These days I roamed around our PyTorch Forums and tried to find a way to initialize the weight matrix. And I found several ways to achieve that. May I ask which one would you recommend most?</p>
<hr>
<p>Suppose we have a very simple (but typical) neural network. And our target is to initialize the weight in the first <code>conv1</code> layer as <code>[[0.,0.,0.],[1.,1.,1.],[2.,2.,2.]]</code>. (As a <code>3*3</code> filter)</p>
<pre><code class="lang-auto">class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.pool = nn.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))

        return x
</code></pre>
<hr>
<p>So here are serval ways that we can initialize the weights: (Huge respect to <a href="https://discuss.pytorch.org/u/vmirly1">vmirly1</a>, <a href="https://discuss.pytorch.org/u/ptrblck">ptrblck</a>, et al.)</p>
<ul>
<li>
<a href="https://discuss.pytorch.org/t/custom-weight-initialization/20544/12">Method 1</a> Define the customize weight matrix inside the <code>__init__</code>:</li>
</ul>
<pre><code class="lang-auto">class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 3, 3)
        self.pool = nn.MaxPool2d(2, 2)

        K = torch.tensor([[0.,0.,0.],[1.,1.,1.],[2.,2.,2.]])  # add the weight here
        K = torch.unsqueeze(torch.unsqueeze(K,0),0) # assign it to the cov1
        self.conv1.weight.data = self.conv1.weight.data * 0 + K

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))

        return x
</code></pre>
<hr>
<ul>
<li>
<a href="https://discuss.pytorch.org/t/custom-weight-initialization/20544/17">Method 2</a>: Define the weight after you build the instance (before your training, of course)</li>
</ul>
<pre><code class="lang-auto">net = Net()

# then change the weights outside the class
K = torch.tensor([[0.,0.,0.],[1.,1.,1.],[2.,2.,2.]])  
K = torch.unsqueeze(torch.unsqueeze(K,0),0) 
net.conv1.weight.data = net.conv1[0].weight.data *0 + K 
</code></pre>
<hr>
<ul>
<li>
<a href="https://discuss.pytorch.org/t/how-to-assign-an-arbitrary-tensor-to-models-parameter/44082/3">Method 3</a> Use the <code>saved-state-dict</code> to update the weights</li>
</ul>
<hr>
<ul>
<li>
<a href="https://pytorch.org/tutorials/beginner/transformer_tutorial.html" rel="noopener nofollow ugc">Method 4</a> Use a class method to achieve that (from tutorials)</li>
</ul>
<pre><code class="lang-auto">class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 3, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.init_weights()
    
    def init_weights(self):
        K = torch.tensor([[0.,0.,0.],[1.,1.,1.],[2.,2.,2.]])
        K = torch.unsqueeze(torch.unsqueeze(K,0),0)
        self.conv1.weight.data = K

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        return x
</code></pre>
<hr>
<p>I think these are all the available methods that I can find on the Internet. May I ask which one is most recommended? Or any one of them is risky?</p>
<p>Thank you in advance!</p>

Superklez: <p>I personally prefer method 4 because it can be really convenient when, for example, you want to initialize the weights of multiple layers. Here’s one example of when it can be helpful:</p>
<pre><code class="lang-auto">    def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')

                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

            elif isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')

                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
</code></pre>

ptrblck: <p>I also like <a class="mention" href="/u/superklez">@Superklez</a>’s approach for the same mentioned reason. In case you just want to assign known values to a single layer, I would probably use “Method 1”.</p>
<p>While your approaches would work fine, I would not recommend to use the <code>.data</code> attribute in any of them, as it might yield unwanted side effects. You could assign a new <code>nn.Parameter</code> to the <code>weight</code> attribute directly (and by wrapping it into a <code>with torch.no_grad()</code> block if necessary), use the <code>nn.init</code> methods as seen in <a class="mention" href="/u/superklez">@Superklez</a>’s code, or the <code>.copy_</code> method in case you want to assign the values directly to a parameter.</p>

Nick_ishere: <p>Thank you for your suggestion, <a href="https://discuss.pytorch.org/u/Superklez">Superklez</a>. Your example greatly explains how to use control-flow to initialize distinct layers. Really helpful!</p>

Nick_ishere: <p>Hi <a href="https://discuss.pytorch.org/u/ptrblck">ptrblck</a>, thank you a lot for your explanation. I just test the method you mentioned, it works great as follows:</p>
<pre><code class="lang-auto">class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 3, 3)
        self.pool = nn.MaxPool2d(2, 2)

        # now assign the parameters
        K = torch.tensor([[0.,0.,0.],[1.,1.,1.],[2.,2.,2.]])
        K = torch.unsqueeze(torch.unsqueeze(K,0),0)
        
        with torch.no_grad():
            self.conv1.weight = nn.Parameter(K)
    
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        return x

net = Net()
</code></pre>
<hr>
<p>However, I am quite confused with <code>torch.no_grad()</code> here. As it is mentioned <a href="https://pytorch.org/docs/stable/generated/torch.no_grad.html" rel="noopener nofollow ugc">in the document</a>, <code>with torch.no_grad()</code> could set the <code>require gradient</code> into <strong>False</strong> and stop the gradient calculation. Nevertheless, when I print out the parameters, all of them <strong>are still trainable</strong>. May I ask why this happen?</p>
<pre><code class="lang-auto">for param in net.parameters():
    print(param)

Parameter containing:
tensor([[[[0., 0., 0.],
          [1., 1., 1.],
          [2., 2., 2.]]]], requires_grad=True)
Parameter containing:
tensor([ 0.0328, -0.2026, -0.2669], requires_grad=True)
</code></pre>

ptrblck: <p><code>torch.no_grad()</code> will make sure that the operations inside the block are not tracked by Autograd and thus not recorded in the computation graph (as you don’t want to backpropagate through the parameter assignment).<br>
The <code>nn.Parameter</code> itself should keep its <code>requires_grad=True</code> attribute.</p>

Nick_ishere: <p>Aha Got it! Thank you so much for the explanation, ptrblck. Really helpful! <img src="https://discuss.pytorch.org/images/emoji/apple/clap.png?v=9" title=":clap:" class="emoji" alt=":clap:"> <img src="https://discuss.pytorch.org/images/emoji/apple/clap.png?v=9" title=":clap:" class="emoji" alt=":clap:"></p>