Custom weight initialization in PyTorch

Question

What would be the right way to implement a custom weight initialization method in PyTorch?

I believe I can't directly add any method to 'torch.nn.init` but wish to initialize my model's weights with my own proprietary method.

Answer 1

You can define a method to initialize the weights according to each layer:

def weights_init(m):
    classname = m.__class__.__name__

    if classname.find('Conv2d') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

And then just apply it to your network:

model = create_your_model()
model.apply(weights_init)

Answer 2

See https://discuss.pytorch.org/t/how-to-initialize-weights-bias-of-rnn-lstm-gru/2879/2 for reference.

You can do

weight_dict = net.state_dict()
new_weight_dict = {}
for param_key in state_dict:
     # custom initialization in new_weight_dict,
     # You can initialize partially i.e only some of the variables and let others stay as it is
weight_dict.update(new_weight_dict)
net.load_state_dict(new_weight_dict)