Custom weight initialization in PyTorch

问题

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.

回答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)

回答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)

来源：https://stackoverflow.com/questions/51125856/custom-weight-initialization-in-pytorch