모델 및 레이어 가중치 초기화(weights initialization)

이번에 논문 구조를 구현하는 과정에서 (DCGAN, LSGAN) 가중치 초기화하는 과정을 거쳤는데요

어이없는 착각을 하면서 며칠간 가중치 초기화만 만졌습니다..

그래서 가중치 초기화 이 녀석이랑 너무 친해진 거 같아서 다른 분들께 도움이 되고자 합니다.

 

1. Why 가중치 초기화?

왜 하는지에 대해서는 이미 다른 분들이 많이 포스팅하였더라고요.

gradient 관련해서 많이들 쓰시는데 쓰시는 Neural Network에 맞게 쓰시면 될 거 같습니다
(He  / Xavier)

그리고 논문에서 모델 구조에 대해 설명하면서 가중치 초기화에 대해 언급하는 경우가 있으니

그럴 경우 그에 맞게 초기화하면 될 거 같습니다.

 

2. 가중치 초기화 코드

net = nn.Sequential(
    nn.Conv2d(3, 2, 3), nn.Linear(2, 3)
)

일단 가벼운 network를 준비합니다.

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
        print('layer success')
    elif classname.find("BatchNorm") != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)
        print('batch success')

첫번째 방법은 class의 이름을 활용해서 초기화 함수를 정의한 것입니다.

장점은 ConvTranspose2d와 Conv2d 모두 'Conv' 가 들어가기 때문에 같이 적용시켜줄 수 있습니다.

def weights_init(m):
  if type(m) == nn.Linear:
    torch.nn.init.xavier_normal_(m.weight)
    torch.nn.init.zeros_(m.bias)
  elif isinstance(nn, nn.Conv2d):
    torch.nn.init.xavier_uniform_(m.weight)
    torch.nn.init.zeros_(nn.bias)

두번째 방법은 이렇게 class이름이 아닌 type으로 정의해주는 방법입니다.

이 방법을 쓰실 경우 ConvTranspose2d 를 if문에 추가해주셔야합니다.

net.apply(weights_init_normal).state_dict()

그리고 초기화 함수를 적용해주기 위해 apply함수를 사용하는데요,

여기서 apply함수를 쓰면은 해당 모듈의 sub module들(nn.Conv2d , nn.Linear)에게 해당 함수가(가중치 함수)

적용이 됩니다. net안에 있는 레이어들에게 모두 적용이 되었는지 확인해보겠습니다.

 

layer success
OrderedDict([('0.weight', tensor([[[[ 1.1839e-02,  1.2020e-02, -4.4393e-02],
                        [ 1.6813e-02, -1.2994e-02,  2.7541e-02],
                        [ 1.3415e-02, -1.4987e-02,  1.9120e-02]],
              
                       [[ 2.3735e-02, -3.8407e-03, -1.8328e-02],
                        [ 9.5370e-03, -6.9400e-05,  5.7889e-03],
                        [-3.6280e-02,  2.2017e-02,  1.6323e-03]],
              
                       [[-3.7677e-03,  5.7425e-02,  4.0497e-02],
                        [-7.9902e-03, -4.4012e-03,  3.2745e-02],
                        [-5.8548e-03,  1.5467e-02,  7.2594e-03]]],
              
              
                      [[[ 2.7921e-03,  2.4599e-02,  3.7741e-02],
                        [-2.5849e-02, -2.1165e-02,  2.2447e-03],
                        [-1.4344e-02,  1.5034e-02,  1.0269e-02]],
              
                       [[-9.1560e-03, -2.4378e-02, -4.2211e-03],
                        [-2.3763e-03, -1.0574e-02,  2.2381e-02],
                        [-1.1711e-02, -1.4624e-03, -4.7581e-02]],
              
                       [[-1.9173e-02, -8.9866e-03, -8.9565e-03],
                        [-2.2351e-02,  1.8486e-02,  2.5473e-02],
                        [ 1.3440e-02, -2.3967e-02,  1.2505e-02]]]])),
             ('0.bias', tensor([-0.1627, -0.0732])),
             ('1.weight', tensor([[-0.0766, -0.0315],
                      [-0.5720,  0.5564],
                      [-0.2420,  0.3670]])),
             ('1.bias', tensor([ 0.6075, -0.3196, -0.1791]))])

다 작은수로 초기화가 되었음을 확인할 수 있습니다.

그렇다면 이렇게 간단한걸 왜 제가 며칠동안 고생했을까요?(페북 페이지에도 물어보고 난리가 아니었습니다)

예를 들어 -8.9866e-03 이 값을 그냥 -8로 본겁니다..

하하 정말 바보 같죠? 

그래서 초기화를 해주었는데 저런 값들을 보고 가중치 초기화가 되지 않았다고 착각했네요..

그래도 덕분에 며칠간 가중치 초기화만 만져서 더 친해진 기분입니다^.^

(DCGAN의 결과가 잘 나오지 않길래 가중치 초기화가 잘 안되서 그런줄 알았는데 아니었네요.. 다시 한번 코드를 만져보겠습니다..)

 

그리고 마지막으로 함수를 사용하지 않고 초기화하는 방법을 소개해드리겠습니다.

for layer in generator.modules():
  if isinstance(layer, nn.ConvTranspose2d):
    layer.weight.data.normal_(0.0, 0.02)
    print('성공')

 

근데 함수를 써서 초기화 해주시는 방법을 추천드립니다!

 

저처럼 실수해서 삽질하는 시간 없었으면 좋겠습니다.