PGGAN을 이용한 강아지, 고양이 생성모델 만들기 (2)

이번 글에서는 지난 글에서 만든 데이터셋과 모델을 이용하여 훈련을 시켜보도록 하겠습니다.

 

먼저 데이터셋과 model state dict를 가져오기 편리하게 steps별로 path를 리스트형태로 정의해줍니다.

steps는 모델의 steps와 같은 의미입니다.

test_size는 train log에 찍을 이미지의 개수입니다.

 

dataset의 directory_path는 미리 정의해둔 dataset_path의 self.steps번째에서 찾을 수 있습니다.

dataset을 로드하여 trainloader와 validloader를 정의해줍니다.

생성자와 판별자도 정의해주고요. loss function과 optimizer도 정의해줍니다. test_z는 train log에 사용할 잠재벡터입니다. 마지막으로 모델의 state_dict를 로드합니다.

model의 state_dict는 한 블록마다 따로따로 저장해줍니다.

 

load model 메소드를 좀 더 자세히 보도록 하겠습니다.

 

generator는 step이 하나씩 늘어날 때마다 끝에 블록이 한개씩 추가되기 때문에,

0 ~ 현재 층수 -1까지 그 전 steps의 state dict를 사용하면 됩니다.

 

discriminator는 step이 하나씩 늘어날 때마다 앞에 블록이 한개씩 추가되기 때문에,

1 ~ 현재 층수까지 그 전 steps의 state dict를 사용하면 됩니다.

 

clear cuda memory 메소드는 cuda memory를 비워주는 메소드입니다.

test 메소드는 이전 글에서 만든 merge_test_pred 메소드를 이용하여 train_log를 찍는 메소드입니다.

각 steps에 해당하는 resolution의 디렉토리에 저장해줍니다.

 

train 메소드는 generator와 discriminator를 학습하는 메소드입니다.

alpha값이 선형적으로 증가될 수 있도록 매 반복마다 alpha_gap을 더해줍니다.

 

valid 메소드는 valid loss를 확인하기 위한 메소드입니다.

 

 

run 메소드는 실제 학습을 진행하는 메소드입니다.

train history와 valid history에 한번의 epoch마다 generator_loss와 discriminator_loss를 추가합니다.

alpha의 초기값은 0이고, alpha_gap은 1 / (한번의 epoch에서 반복하는 횟수 * epochs)입니다. 이렇게 alpha_gap을 정의해준다면, 학습이 끝나는 순간의 alpha값은 1이 될 것입니다.

 

그 다음은 모델들을 훈련, 검증하며 loss를 출력합니다.

한번의 epoch의 끝날때마다 test 메소드를 이용하여 train_log를 찍어줍니다.

마지막으로 train_history와 valid_history를 반환해줍니다.

 

훈련을 하기위한 모든 준비는 끝이 났습니다.

이제 이렇게 훈련을 하면 모두 끝입니다!

 

제 훈련환경은 Google Colab Pro입니다.

 

alpha가 1이 되어야 다음 steps로 넘어갈 수 있기 때문에

중간에 백업해둔 state dict를 갖고 그 다음 steps로 넘어갈 수 없어서 훈련이 더 힘들었습니다.

 

예를 들어 100epochs에서 학습을 하고 있는데, 80epoch에서는 모델 성능이 굉장히 좋아보이는데, 100epoch에서 모델 성능이 많이 떨어진 것 같다면 보통의 경우에는 80epoch에서의 state dict를 사용할텐데, PGGAN 구조상 80epoch에서의 alpha값은 1이 아니기 때문에 그러지 못합니다.

 

마지막 epoch에서 성능이 좋아야만 다음 steps로 넘어갈 수 있어서 각 resolution마다 여러번 학습을 하였기 때문에 학습시간이 총 4~5일 정도 걸렸습니다.

 

결과는 다음과 같습니다.

PGGAN 이론 글에서는 Minibatch Standard Deviation 기법으로 모드 붕괴를 어느 정도 완화한다고 했지만, 데이터셋의 다양성도 충분하여야 한다는 전제조건이 있었습니다. 비교적 데이터셋의 크기가 작았기 때문에 생긴 문제라고 생각합니다. 그렇지만 모드붕괴 현상이 아주 심각해보이지는 않습니다. 적은 데이터셋으로 낸 결과로 납득할만한 성능입니다.

 

train log를 gif 파일로 만드는 코드는 위와 같습니다.

 

그리고 그 결과입니다.

https://www.youtube.com/watch?v=pvSaE_BVKJM 

 

train.py 코드 - https://github.com/SimplePro/PGGAN/blob/main/train.py

GitHub - SimplePro/PGGAN: make cat and dog images using PGGAN

github repository - https://github.com/SimplePro/PGGAN

GitHub - SimplePro/PGGAN: make cat and dog images using PGGAN

 

이렇게 PGGAN을 이용한 강아지, 고양이 생성모델 만들기 챕터가 끝났습니다.

며칠간 이번 프로젝트를 진행하면서 새로 배운 점이 많았고, 그 전보다 성장한 것 같아 뿌듯하네요.

 

감사합니다.