[Pytorch] 파이토치 가중치 초기화 방법(Xavier, He)

Python torch Weight Initialization

파이토치에서 Xavier, He 등의 가중치 초기화를 각 layer 별로

혹은 모델 전체에 대하여 진행하는 방법을 간략하게 요약해보도록 하겠습니다.

 

우선, 다음과 같이 fc1, fc2, fc3의 3개의 layer를 가진 간단한 모델 구조를 가정해 보겠습니다.

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.fc1 = nn.Linear(8, 4)
        self.fc2 = nn.Linear(4, 2)
        self.fc3 = nn.Linear(2, 1)

    def forward(self, x):
        return self.fc3(self.fc2(self.fc1(x)))

model = Net()

 

 

Xavier 가중치 초기화

대표적인 Xavier 가중치 초기화는 uniform, normal 분포의 두 가지 버전이

흔하게 사용되는데 torch.nn에서 두 가지 초기화 메소드를 모두 제공하여

쉽고 편리하게 적용할 수 있습니다. 각 버전에 대한 수식은 다음과 같습니다.

 

1. Xavier Uniform Initialization

$$ W \sim U(-\sqrt{ \frac{6}{n_{in} + n_{out}}}, \sqrt{ \frac{6}{n_{in} + n_{out}}}) $$

 

2. Xavier Normal Initialization

$$ W \sim N(0, Var), \; Var = \sqrt{ \frac{2}{n_{in} + n_{out}}} $$

 

두 가지 버전의 가중치 초기화를 파이토치에서 적용하는 방법은 아래와 같습니다.

먼저, fc1 layer 1개에 대하여 따로 적용한 예시입니다.

# fc1 layer에 대한 xavier uniform initialization
nn.init.xavier_uniform_(model.fc1.weight)

# fc1 layer에 대한 xavier normal initialization
nn.init.xavier_normal_(model.fc1.weight)

 

모델 내 layer 전체 혹은 layer 이름을 기준으로 지정하여 가중치 초기화를 진행하는 방법은

다음 예제의 코드와 같습니다.

# 전체 layer에 대하여 xavier uniform initialization 예시
for name, child in model.named_children():
    nn.init.xavier_uniform_(child.weight)

# fc2, fc3 layer에 대하여 xavier normal initialization 예시
for name, child in model.named_children():
    for param in child.parameters():
        if name in ['fc2', 'fc3']: # 원하는 layer 이름 지정
            nn.init.xavier_normal_(child.weight)

 

반응형

 

He 가중치 초기화

활성화 함수로 ReLU 계열을 사용하는 경우에는 Xavier 가중치 초기화를 할 경우

값들이 0으로 수렴하는 현상이 나타날 수 있어 He 초기화 방법을

대신해서 주로 사용합니다. 해당 방법의 수식은 다음과 같습니다.

 

1. He Uniform Initialization

$$ W \sim U(-\sqrt{ \frac{6}{n_{in}}}, \sqrt{ \frac{6}{n_{in}}}) $$

 

2. He Normal Initialization

$$ W \sim N(0, Var), \; Var = \sqrt{ \frac{2}{n_{in}}} $$

 

torch.nn에서는 He 가중치 초기화에 대한 메소드도 마찬가지로 제공하는데,

kaiming_uniform_, kaiming_normal_의 이름으로 제공하고 있습니다.

 

fc1 layer 1개에 대하여 따로 적용한 예시는 아래와 같습니다.

# fc1 layer에 대한 He uniform initialization
nn.init.kaiming_uniform_(model.fc1.weight)

# fc1 layer에 대한 He normal initialization
nn.init.kaiming_normal_(model.fc1.weight)

 

마찬가지로 모델 내 layer 전체 혹은 layer 이름을 기준으로 지정하여

He 가중치 초기화를 진행하는 방법도 위의 Xavier의 예시에서

메소드명만 바꿔주시면 동일하게 적용됩니다.

 

 

기타 가중치 초기화 방법

이 외에도 torch.nn에서는 균일 분포, 정규 분포, 상수, 단위 행렬 등

다양한 방법의 가중치 초기화 방법들을 제공하고 있으니,

이에 대해서는 아래 링크의 파이토치 공식 document를 참고해주시면 됩니다.

https://pytorch.org/docs/stable/nn.init.html

torch.nn.init — PyTorch 1.13 documentation

 

여기까지 파이토치의 가중치 초기화 방법에 대한 예제를 다룬 글을 마치겠습니다. 감사합니다.