Пользовательская инициализация веса в PyTorch

Question

Как правильно реализовать метод custom weight initialization в PyTorch?

Я полагаю, что не могу напрямую добавить какой-либо метод в torch.nn.init, но хочу инициализировать вес моей модели с помощью собственного проприетарного метода.

Answer 1

Вы можете определить метод для инициализации весов в соответствии с каждым слоем:

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)

А затем просто примените его к своей сети:

model = create_your_model()
model.apply(weights_init)

Answer 2

См. https://discuss.pytorch.org/t/how-to-initialize-weights-bias-of-rnn-lstm-gru/2879/2 для справки.

Вы можете сделать

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)