StyleGAN2-ADA 학습 튜토리얼 (AHFQ-v2 Dataset 사용)

AHFQ-v2 Dataset을 통해서 StyleGAN2-ADA를 학습하는 방법을 정리하려고 한다. 

이 글을 참고해서 dataset만 바꿔주면 custom dataset으로도 충분히 학습을 진행할 수 있을 것이다. 

 

Dataset 준비

나는 동물 얼굴 데이터셋인 ahfq-v2 데이터셋을 사용하였다. 

다음 링크에 접속하면 

stargan-v2/README.md at master · clovaai/stargan-v2 (github.com)

stargan-v2/README.md at master · clovaai/stargan-v2

 

다음과 같이 dataset 설명과 다운로드할 수 있는 방법이 나온다. 

 

여기서 download.sh 파일만 다운로드한 후, 다음 명령어를 실행해 주면 해당 데이터셋을 다운로드할 수 있다. 

다운로드한 후에는 zip 파일을 원하는 경로로 옮겨주면 된다. custom dataset를 사용하고 싶다면 custom dataset를 data 경로로 옮겨준다. 

bash download.sh afhq-dataset

 

이때 나는 zip파일 그대로를 다운로드하고 싶었기 때문에 sh 파일에서 afhq-v2-dataset 다운로드하는 코드에서 unzip과 rm 가 들어있는 코드를 주석 처리해서 사용하였다. 

elif  [ $FILE == "afhq-v2-dataset" ]; then
    #URL=https://www.dropbox.com/s/scckftx13grwmiv/afhq_v2.zip?dl=0
    URL=https://www.dropbox.com/s/vkzjokiwof5h8w6/afhq_v2.zip?dl=0
    ZIP_FILE=./data/afhq_v2.zip
    mkdir -p ./data
    wget -N $URL -O $ZIP_FILE
    # unzip $ZIP_FILE -d ./data
    # rm $ZIP_FILE

 

훈련 시작

훈련을 시작하는 방법은 간단하다.

우선, 원하는 환경을 구축한다. 

Dockerfile을 제공하기 때문에 docker 환경 내에서 실행해도 되고, 본인은 anaconda 가상환경을 만들어서 실행하였다. 

환경을 구축했으면 환경 내부에서 stylegan2-ada 코드를 clone 한다. 

git clone https://github.com/NVlabs/stylegan2-ada.git

 

stylegan2-ada 폴더로 이동한 후 필요한 라이브러리를 설치한다. 

pip install click requests tqdm pyspng ninja imageio-ffmpeg==0.4.3

 

다음 명령어로 훈련을 시작한다. 

python train.py --outdir=./training-runs --data=./datasets/afhq_v2.zip --gpus=1

경로, gpu 개수 등은 본인의 디렉토리에 맞게 설정해 주면 된다.

주요 옵션은 다음과 같고 꼭 설정해줘야 하는 옵션은 볼드체로 작성해 놓았다. (--outdir, --data)

• --outdir: 결과 저장 폴더
• --gpus: 사용할 gpu 개수 (default: 1)
• --snap: snapshot interval (default: 50 ticks)
• --seed: random seed (default: 0)
• --data: 훈련 데이터 경로 (directory or zip)
• --cond: 데이터셋 라벨을 기반으로 conditional model 학습할지 (default: false)
• --cfg: config (default: auto) (옵션: 'auto', 'stylegan2', 'paper256', 'paper512', 'paper1024', 'cifar’)
  • 'auto': dict(ref_gpus=-1, kimg=25000, mb=-1, mbstd=-1, fmaps=-1, lrate=-1, gamma=-1, ema=-1, ramp=0.05, map=2)
• --kimg: training duration 설정
• --batch: batch size 결정
• --aug: augmentation mode 설정 (default: ada) (옵션: ‘noaug’, ‘ada’, ‘fixed’)
• --resume: 이어서 학습할 경우 pkl 경로 설정

 

혹시 중간에 이런 오류가 발생한다면 다음 링크를 참고해서 해결하면 된다.

RuntimeError: derivative for aten::grid_sampler_2d_backward is not implemented

[오류 해결] RuntimeError: derivative for aten::grid_sampler_2d_backward is not implemented(feat. StyleGAN2-ADA) — still_growing (tistory.com)

[오류 해결] RuntimeError: derivative for aten::grid_sampler_2d_backward is not implemented(feat. StyleGAN2-ADA)

 

 

중간 결과

본인의 경우 500 tick(2000 kimg)까지 학습하는 데에  RTX 3090 기준 1일 22시간 28분가량 소요되었다. 기본적으로 25000 kimg로 설정되어 있기 때문에 기본 설정인 상태로 끝까지 돌린다면 24일이 넘게 걸릴 것으로 예상된다.

기본적으로 100 tick(400 kimg)마다 이미지와 pkl 파일이 저장되는데 중간 결과를 공유하자면 다음과 같다. 

초반에는 조금 뭉개진 얼굴이 많이 보이다가 훈련이 거듭될수록 점점 그럴듯한 동물 얼굴 이미지가 생성되는 것을 확인할 수 있다. 

 

실제 이미지 모음

 

fakes000000(kimg 기준)

 

fakes000200

동물들 얼굴이 섞인 이미지가 많이 보이고 얼굴이 많이 뭉개진 형태가 많이 보인다.

 

fakes000400

 

fakes000800

 

fakes001200

 

fakes001600

 

fakes002000

 

한 이미지만 비교했을 때

 

Inference - 이미지 생성

다음 이미지를 타겟 이미지로 훈련시킨 다음 2000 kimg까지 훈련시켰을 때의 network를 이용해서 이미지를 생성해 보았다. 

1. 이미지 projection

다음 명령어를 실행하면 주어진 이미지 파일에 대해서 일치하는 latent space를 찾을 수 있다. 

경로는 본인 경로에 맞게 설정해 주면 된다.

python projector.py --outdir=out --target=pixabay_dog_000135.png --network=./training-runs/00002-afhq_v2-auto1/network-snapshot-002000.pkl

위 명령어에 대한 출력은 다음과 같고 

• target.png : 타겟으로 설정했던 이미지
• proj.png : latent space를 찾고 생성한 이미지
• projected_w.npz : latent space
• proj.mp4 : latent space를 찾는 과정을 담은 비디오

proj.mp4과 proj.png 는 다음과 같이 생성되었다. 

proj.mp4

 

proj.png

 

imageio관련해서 아래 오류가 날 수도 있는데 다음 링크를 참고해서 해결하면 된다.

valueerror: could not find a backend to open 'out/proj.mp4' with iomode 'wi'. based on the extension, the following plugins might add capable backends: ffmpeg: pip install imageio[ffmpeg] pyav: pip install imageio[pyav]

https://seoy00.tistory.com/84

valueerror: could not find a backend to open `out/proj.mp4`` with iomode `wi`. based on the extension, the following plugins mig

 

 

2. projection 된 latent vector를 통해 이미지 생성하기

다음 명령어를 통해 이미지를 생성할 수 있는데 이때 --projected-w를 통해서 이미지 생성에 사용할 latent space를 결정할 수 있다. 

위에서 생성된 latent vector인 npz를 가지고 이미지를 생성하면 당연히 위의 이미지와 동일한 이미지가 생성된다. 

python generate.py --outdir=out --projected-w=out/projected_w.npz --network=./training-runs/00002-afhq_v2-auto1/network-snapshot-002000.pkl

생성된 이미지는 다음과 같다. 

proj00.png

 

단일 이미지 생성하기

latent space를 뽑지 않고 seed 값을 조정해서 그냥 데이터를 생성하는 방법도 있다. 다음 명령어를 사용하면 된다. 

python generate.py --outdir=out --trunc=0.7 --seeds=0 --network=./training-runs/00002-afhq_v2-auto1/network-snapshot-002000.pkl

옵션 설명은 다음과 같다. 

• --outdir: 이미지 저장할 경로를 설정
• --network: 생성하고 싶은 시점의 pkl 파일 경로를 설정
• --seeds: 모델에서 무작위 seed 선택
• --truncation: 숫자가 작을수록 더 사실적으로 표현되고 이 변수는 다양성에도 영향을 미친다. 값은 무한대까지 선택할 수 있지만 대부분은 0.5와 1.0 중 하나를 선택한다고 한다.

대부분은 0.5와 1.0 중 하나를 선택한다고 하는데 그 이유가 궁금해서 seed값을 고정한 상태에서 truncation만 변화를 주어 이미지를 생성해보았다. 

(왼쪽 위에서부터

0.1, 0.2, 0.3, 0.4, 0.5,

0.6, 0.7, 0.8, 0.9, 1.0

1.5, 2.0, 10)

stylegan2-ada에서는 truncation값이 낮으면(예 0.1) 생성된 이미지가 평균적이고 안정적인 특성을 지니기 때문에 생성된 이미지들이 대부분 평균에 가까운 모양을 가지며 노이즈가 적고 일된 특징을 보여준다. 하지만, 최대한 안정적으로, 즉 모든 이미지의 평균에 가까운 형태를 띄기 때문에 나타내고자 하는 이미지에서 거리가 있어 보인다. 

Truncation 값이 중간일 때는 (예: 0.5, 0.7)에는 모델이 생성하는 이미지가 조금 더 다양해지며, 현실감 있는 변화를 나타내고 안정적인 것을 확인할 수 있다. 확실히 왜 0.5에서 1사이의 값으로 이미지를 생성하는지 알 것 같다. 

Truncation 값이 1 이상으로 커질 때는 모델이 생성하는 이미지가 평균에서 멀어지면서 더 독창적이거나 때로는 매우 왜곡된 형태로 나타날 수 있다고 한다. 여기서는 Truncation 값이 10까지 설정되어 있어, 이로 인해 마지막 이미지에서 매우 추상적이고 혼란스러운 패턴이 생성된 것을 확인할 수 있다. 

 

Reference

https://github.com/NVlabs/stylegan2-ada-pytorch