PyTorch Custom Weight Initialization Example

Neural network weight and bias initialization is a surprisingly tricky topic. In most cases, the default initialization code works fine, even though it’s poorly documented and the source code is extremely complex.

Just to test my understanding of weight and bias initialization I decided to put together a demo. First, I created a 6-(10-10)-3 neural network where I used the built-in torch.nn.init functions. I trained the network and recorded the results. The network had poor accuracy but that wasn’t the point.


Left: A neural network trained using the built-in torch.nn.init initialization functions. Right: The same network but initialized using custom code. The output of the two approaches is identical.

Next, I edited my program and used custom code to initialize weights and biases to the same values as supplied by the built-in torch.nn.init functions. I trained the network and observed that the results were identical to the first version. Therefore I was pretty sure I understood PyTorch weight initialization.

My demo program predicts employee job-type from sex, age, city, and income. The data is synthetic.

Neural network using built-in torch.nn.init functions:

class Net(T.nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    self.hid1 = T.nn.Linear(6, 10)  # 6-(10-10)-3
    self.hid2 = T.nn.Linear(10, 10)
    self.oupt = T.nn.Linear(10, 3)

    version 1: built-in init functions
    T.nn.init.constant_(self.hid1.weight, 0.1)
    T.nn.init.zeros_(self.hid1.bias)
    T.nn.init.constant_(self.hid2.weight, 0.2)
    T.nn.init.zeros_(self.hid2.bias)
    T.nn.init.constant_(self.oupt.weight, 0.3)
    T.nn.init.zeros_(self.oupt.bias)

  def forward(self, x):
    z = T.tanh(self.hid1(x))
    z = T.tanh(self.hid2(z))
    z = self.oupt(z)  # no softmax: CrossEntropyLoss() 
    return z

Neural network using custom weight and bias initialization:

class Net(T.nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    self.hid1 = T.nn.Linear(6, 10)  # 6-(10-10)-3
    self.hid2 = T.nn.Linear(10, 10)
    self.oupt = T.nn.Linear(10, 3)

    # version 2: init from scratch
    for i in range(6):
      for j in range(10):
        self.hid1.weight.data[j][i] = 0.1
    for j in range(10):
      self.hid1.bias.data[j] = 0.0

    for i in range(10):
      for j in range(10):
        self.hid2.weight.data[j][i] = 0.2
    for j in range(10):
      self.hid2.bias.data[j] = 0.0

    for i in range(10):
      for j in range(3):
        self.oupt.weight.data[j][i] = 0.3
    for j in range(3):
      self.oupt.bias.data[j] = 0.0

  def forward(self, x):
    z = T.tanh(self.hid1(x))
    z = T.tanh(self.hid2(z))
    z = self.oupt(z)  # no softmax: CrossEntropyLoss() 
    return z

The moral of the story is that custom weight and bias initialization in PyTorch is possible, but it’s not necessary in any normal scenario I can think of.



When I was younger and had more time, I loved to read fiction novels. Here are three initial books in series that I like a lot. Left: “Richard Bolitho – Midshipman” (1975) by Alexander Kent. Center: “Casino Royale” (1953) by Ian Fleming. Right: “A Study in Scarlet” (1887) by Arthur Conan Doyle.


This entry was posted in PyTorch. Bookmark the permalink.