지금 Github에는 InterFaceGAN에 대한 코드가 다 올라와있지만, 아직 이를 적용할 수 있는 모델은 PGGAN과 StyleGAN 뿐이다. 물론 Issue에 StyleGAN2를 해달라는 요청과 지금 진행중이라는 답변이 있긴한데 대략 1년이 지난 것을 보아 아마 까먹고 다른 연구에 몰두 중이신 것 같다. 아무튼 이제 내가 써볼 모델은 StyleGAN2이고 보유중인 pretrained weight 역시 StyleGAN2 base라 어쩔 수 없이 자체 포팅을 진행해보았다.

Model Architecture Analysis

전반적인 모델의 구조를 보았을 때는 다음과 같이 진행되는 것 같다.

base_generator에서는 InterFaceGAN에서 해당 GAN 모델을 사용할 때 필요한 추가적인 함수들이 정의되어 있고 요 코드를 기반으로 각 GAN model의 generator가 정의되어 있었다. 그리고 원래 존재하는 fundamental한 GAN의 Architecture의 정보는 gan_generator_model.py에 저장되어 있었다. 그리고 마지막으로 이러한 모델들의 metadata를 담아놓은 config 느낌에 model_settings.py 파일이 존재했다.

StyleGAN 2 Porting

그렇기에 우선 stylegan2 model architecture를 sefa에 있던 코드를 불러와 stylegan2_generator_model.py에 저장을 하고 stylegan_generator를 그대로 복사해 stylegan2_generator.py를 만들어 주었다.

그 후 stylegan과 stylegan2의 각 model파일과 generator 파일을 펼쳐놓고 서로 다른 부분에 대해서 맞춰가기 시작했다.다행히도 큰 구조 자체에는 다른 부분이 없어서 사용 변수들만 형식에 맞게 수정해주면 되었다. 우선 stylegan2에서만 존재했지만 진행 상황을 표시하기 위한 변수들을 삭제해주었다. 그 외에는 truncation 관련 hyperparameter를 어떻게 주고받는지에도 차이가 없어서 해당 부분을 통일시켜 주었다.

이렇게 코드 상으로의 문제점을 1차적으로 잡고 prepare_data라는 dataset을 준비하는 코드를 stylegan2로 세팅하여 돌려보았다. 우선 시작하자마자 문제가 발생했는데, 우리가 필요한 건 단순 Generator지만 저장된 weight파일은 Discriminator도 보유해서 약간의 문제가 발생했다. Generator만 불러올 수 있게 처리를 해주고, stylegan2_model에서는 tensor를 dictionary에 담아서 전달해 주었지만 stylegan에서는 direct하게 전달해주는 것을 확인할 수 있어서 해당 값 중간중간에 다 key값을 삽입해주었다.

우선 ffhq dataset으로 실험을 해보았을 때, 성공적으로 돌아가는 것을 확인할 수 있었다. 근데 원래 존재하는 모델 자체도 ffhq 였기에 혹시 꼬이진 않았을까하고, 존재하지 않던 cat과 dog dataset을 돌려보았다. 사실 해상도 차이가 있어서 처음부터 문제가 나길래 잘 작동하는 것 같긴 했지만 결과는...

짜잔
잘 나온다!

Auxiliary Network

이제는 boundary detection을 해줘야 하는데, 앞선 포스팅에서도 말했듯 InterFaceGAN은 supervised learning이기 때문에 prediction을 진행해 줄 auxiliary network가 필요하다. 하지만... Get your own predictor... 눈물을 머금고 issue를 뒤져봤더니 다행히 StyleGAN 자체에서 제공하는 linear seperablitiy 코드가 존재했다. 하지만 StyleGAN 원래 거는 tensorflow... 소종이었던 것에 대해서 StyleGAN을 tf에서 pytorch로 돌리는데 이틀을 새웠던 악몽이 다시 떠올랐다. 그래서 이번에는 굳이 이 도커에서 진행하는 것이 아니라 전처리의 일부분이라고 생각하고 stylegan 전용 도커파일을 찾아 새로운 도커를 파서 진행을 했다.

확실히 오래된 모델이라 그런지 Google Drive에서 모델을 받아오는 데 문제도 있고 tf dependancy 문제도 존재했다. 드라이브 문제는 모델 직접 받아서 연결시켜주고 tf dependancy는 전용 dockerfile로 해결하고 막상 돌려보았는데 결과를 이거 어떻게 보는지를 모르겠다. 그냥 단순히 로스만 나와서 조금 당황스러운데 생각해보니 이거 dataset에 대한 경로를 지정해 준 기억이 없다. 다시 좀 더 알아봐야겠다...

추가로 StyleGAN2에서 좀 더 사용하기 편하게 제공하는 모델을 발견해서 내일은 이거로 해보거나 아니면 논문에 적혀있는 대로 ResNet-50을 그냥 받아와서 모델 구조 보고 마지막만 binary처럼 짜는 등 코드를 직접 짜서 설계해봐야겠다.

+) 동물 표정도 사람용 predictor가 잘 인식할까? 아니면 인간 predictor...