load_state_dict 성능 문제, [nn.ModuleList]

학습을 끝내고 성능 테스트 때 괜찮은 성능을 보였다.

그런 모델을 저장하고, 다른 코드에서 load_state_dict 으로 불러들여서 사용하려한다.

근데 모델을 저장하기 전 성능이 나오질 않는다.

이럴 땐 일부 레이어를 list 안에 넣어놨는 지 의심해야한다.

model 클래스는 다 nn.Module을 상속해서 정의하는데,

self의 속성으로 레이어를 직접 정의하지 않고,

self.layers = [nn.Linear, nn.Conv2D 등..] 식으로 저장하고 사용했다면

위와 같은 문제가 발생한다.

그 이유는 list로 저장해두면 모델 저장 시, 이 레이어들이 인식되지 않는다.

그래서 list안의 레이어들만 빼놓고 나머지 layer들의 weight만 저장된다.

이에 따라 다른 코드에서 load_state_dict 시, list 안의 layer는 학습되지 않은 날 레이어들이 된다.

이 문제는 nn.ModuleList로 해결 가능하다.

위의 self.layers를

self.layer_module = nn.ModuleList(self.layers)

과 같이 선언하고

self.layer_module을 list처럼 쓰면 말끔히 해결한다. (리스트 처럼 인덱싱 가능, +로 이어 붙이기도 가능)

이걸 몰라서 정말 많은 고생을 했다....