pkl 파일을 pt 파일로 바꾸기(StyleGAN2-ADA, Convert pkl to pt)

커스텀 데이터셋을 가지고 StyleGAN2를 학습시킨 경우 기존 .pkl 파일을 .pt로 바꾸어야 하는 경우가 있습니다. 이를 적용해보고 generator로 확인까지 하는 방법을 간단히 소개하고자 합니다.

tensorflow 1.x 버전(이 글은 1.15.5)이 필요합니다.

!git clone https://github.com/dvschultz/stylegan2-ada-pytorch.git

convert 코드 깃허브
해당 깃헙을 클론해옵니다. 공식 레포지토리에도 convert 파일을 제공하지만 제 경우에는 원하는 결과가 나오지 않아서 해당 깃헙으로 작업하였습니다.

#convert
!python /stylegan2-ada-pytorch/export_weights.py /to/input/network.pkl /to/output/network.pt

깃헙내의 export_weights.py 파일을 통해서 pkl파일을 pt로 바꿀 수 있습니다.
그럼 만든 pt파일로 이미지를 만들어 확인해보겠습니다.

!git clone https://github.com/rosinality/stylegan2-pytorch

StyleGAN2-pytorch 깃허브
해당 깃헙을 클론 받아옵니다.

#확인하기, []오류뜨면 --channel_multiplier 1 넣어주세요, Auto2로 학습한 경우 generator.py에 m_mlp를 2로 바꿔주세요
!python generate.py --size 128 --pics 20 --ckpt /to/network.pt 

주석으로 제가 발생했던 오류를 간단하게 요약해서 적었습니다.
여기서 size는 생성할 이미지의 크기, pics는 몇개를 생성할건지, ckpt는 pt파일 경로를 의미합니다.
클론한 stylegan2-pytorch/smaple에 pics 개수 만큼 이미지가 생성되었다면 성공입니다!