StyleGan3 Custom(그림) Data로 학습하기

오늘은 stylegan3로 간단한 훈련을 진행해보겠습니다
 
전체적인 순서는
데이터 준비 (이미지 파일)  →  데이터셋 만들기  →  훈련
이렇게 진행되겠습니다
 
 
 

데이터 준비

일반적으로 데이터를 준비할 때 실사 이미지를 가지고 준비하지만 저는 오늘 그림 데이터를 준비했습니다
하지만 훈련과정은 동일 하니 참고 하시면 될 것 같습니다
 
 

각각 크기가 다른 집 이미지 500장을 준비 했습니다하지만 stylegan 훈련에 필요한 데이터셋에는 동일한 크기의 이미지(256X256, 512X512, 1024X1024)가 필요해서 먼저는 이미지 크기를 변환 합니다 변환은 꿀뷰를 사용했습니다
 
 

 

• 저는 이미지의 크기를 512X512로 하겠습니다 또 여백  붙이기로 사용했지만 꽉차게 늘리기를 사용하셔도 무방합니다 
• 폴더는 따로 생성하셔서 그 경로에 저장하시는게 용이합니다 설정이 완료되면 시작을 눌러주면 순차적으로 변환이 시작됩니다

 
 

 

• 이렇게 크기가 다른 이미지들이 동일한 크기의 이미지로 변경되었다면 이제 데이터셋으로 변경할 준비가 다 됬으니 stylegan3에 필요한 데이터셋으로 만들어주겠습니다
• 먼저는 stylegan3를 설치한 가상환경과 폴더로 들어가주기 위해 anaconda prompt를 실행합니다

cd C:\Users\admin\workspace\stylegan3

conda activate stylegan3

• 혹시나 설치를 안하셨더라면 제 이전 글을 참고하여 설치하시면 쉽게 설치할 수 있습니다
• https://runingcoding.tistory.com/m/123

Stylegan3 window11 설치

 
 
 

데이터셋 만들기

stylgan3 github에 들어가보니 데이터 셋도 여려 종류가 있어서 저도 기존에 사람에 대한 내용 보다도 그림으로 한번 테스트 해보고 싶어서 이번에 그림 이미지로 데이터를 선택해 봤습니다

https://github.com/NVlabs/stylegan3

GitHub - NVlabs/stylegan3: Official PyTorch implementation of StyleGAN3

 
 

• 이미지 크기를 변환 (정사각형 : 1024, 512, 256)
• FFHQ(https://github.com/NVlabs/ffhq-dataset) : 사람

GitHub - NVlabs/ffhq-dataset: Flickr-Faces-HQ Dataset (FFHQ)

python dataset_tool.py --source=/tmp/images1024x1024 --dest=./datasets/ffhq-1024x1024.zip

• MetFaces(https://github.com/NVlabs/metfaces-dataset) : 그림

GitHub - NVlabs/metfaces-dataset

python dataset_tool.py --source=./dataset/origin-512x512 --dest=./dataset/metfaces-1024x1024.zip

• AFHQv2(https://github.com/clovaai/stargan-v2/blob/master/README.md#animal-faces-hq-dataset-afhq) : 동물

GitHub - clovaai/stargan-v2: StarGAN v2 - Official PyTorch Implementation (CVPR 2020)

python dataset_tool.py --source=./dataset/origin-512x512 --dest=./dataset/afhqv2-512x512.zip

•  --source : 데이터셋을 만들 이미지경로
•  --dest : 데이터셋을 저장할 경로

 

 

python dataset_tool.py --source=./dataset/객체별/house512 --dest=./dataset/metfaces-house-512x512.zip

 

• 이렇게 데이터툴 코드를 사용하여 압축파일이 생성 됬다면 데이터셋은 완성입니다

 
 

• 압축파일을 열어보면 00000폴더와 json파일이 생성된 것을 볼 수 있습니다
• 00000폴더를 들어가보면 데이터셋 툴을 실행시킨 이미지들이 들어가 있는 것을 볼 수 있습니다 
• 00000폴더는 1개의 폴더당 1000장 단위로 00001, 00002, ..... 이런식으로 들어갑니다

 
 
 
 

훈련

• 이제 본격적으로 훈련작업을 진행해보겠습니다
• stylegan3의 훈련에는 상당한 gpu 메모리가 필요해서 gpu가 하나라면 알맞은 훈련 세팅을 하는데 오래걸립니다
• 지금은 4090을 사용해서 크게 문제는 없지만 4090을 쓰지 않는 분들도 계실테니 좀 더 낮은 스펙의 컴퓨터로 테스트해보려고 합니다.
• 일단 제가 테스트 해볼 수 있는 컴퓨터 중에 3080 gpu로 세팅한 데스크탑이 있어서 그 컴퓨터 기준으로 설명드리겠습니다

https://github.com/NVlabs/stylegan3/blob/main/docs/configs.md

GitHub - NVlabs/stylegan3: Official PyTorch implementation of StyleGAN3

• 자신에게 맞는 훈련 세팅을 찾고 싶으시면 이 링크로 들어가셔서 찾아보시면 좋은 세팅값을 찾으실 수 있습니다

 

--data : 데이터 경로 --outdir : 결과 경로 --cfg : 모델 --gpus : GPU 갯수 --batch : 크기 --gamma : 정규화 --mirror=1 : 임의의 x-플립으로 데이터 세트 증가 (이미지 수 2배) (수평 대칭인 데이터 세트(예: FFHQ)에 유용하지만 이미지에 눈에 띄는 비대칭 특징(예: 텍스트 또는 문자)이 포함되어 있으면 해로울 수 있습니다 --cbase=16384 : 네트워크 용량 감소 훈련속도↑ 품질↓(저해상도 일수록 차이가 덜 두드러진다) --batch-gpu : GPU당 배치 크기 설정 --snap : 네트워크 스냅샷 (GPU를 적게 사용할 수록 사용해주는 것이 좋다 기본은 200당 한개의 스냅샷) --resume : 훈련 스냅샷부터 실행 --kimg : 총 훈련박복 횟수 (기본은 25000, 5000만해도 적정한 값이 나온다고 함)

• train.py 코드에 필요한 소스들을 몇가지만 정리했습니다 제가 그냥 경험한 대로 적은거라 틀린 내용이 있을 수도 있으니 그 점은 양해부탁드리겠습니다.....

 
 

• 저는 다음과 같이 훈련을 진행하였습니다

python train.py --outdir=./dataset/training-runs --cfg=stylegan3-t --data=./dataset/metfaces-house-512x512.zip --gpus=1 --batch=16 --gamma=5.9 --mirror=1 --kimg=1000 --snap=50

• 처음 훈련에 진입하면 이 화면에서 상당히 오랜 시간동안 멈춰있는 것 같은데 문제가 있으면 알아서 오류랑 같이 취소되니 안심하시고 기다려 주시면 됩니다. 3080 컴퓨터에서 위에 화면 다음 단계로 진입하는데 1시간 조금 넘게 걸렸습니다....

 

• 훈련이 진행 되면 밑에 내가 훈련하고 있는 폴더와 같이 스냅샷 모델과 예시 이미지가 생성되어 내가 어디까지 진행되고 있는지 확인할 수 있습니다 밑에 사진은 5000kimg까지 훈련을 완료한 자료입니다

• 3080gpu로 5000kimg까지 13일 23시간 33분이 걸렸습니다... 상당히 인내가 필요한 작업인 것 같습니다....
• 그래도 스냅샷으로 중간중간 확인할 수 있어서 좋았습니다

• 여기까지 훈련과정은 완료 되었습니다

 
 

테스트

• 이제는 훈련한 모델을 테스트 해보겠습니다
• 훈련한 모델로 생성할 수 있는 방법은 두가지 인데 이미지와 비디오 형태로 볼 수 있습니다

 

 1. 이미지 생성

python gen_images.py --outdir=out --trunc=1 --seeds=0-999 --network=./dataset/training-runs/00018-stylegan3-t-metfaces-house2-512x512-gpus1-batch16-gamma6.6/network-snapshot-005000.pkl

--network : 훈련모델이 있는 경로 --seeds : 생성할 이미지 번호

 

• 위 코드를 실행 시키면 이렇게 이미지를 생성하고 있습니다

 

• 생성이 완료 됬으면 stylegan3 폴더 안에 out 폴더가 만들어지고 들어가보시면 이렇게 새로운 이미지를 생성합니다

 
 
 2. 비디오생성

python gen_video.py --output=house.mp4 --trunc=1 --seeds=0-35 --grid=3x3 --network=./dataset/training-runs/00018-stylegan3-t-metfaces-house2-512x512-gpus1-batch16-gamma6.6/network-snapshot-005000.pkl

• -- grid 는 배열입니다 원하시는 형태로 만드시면 되고 대신 --seeds의 갯수가 배열의 배수의 갯수로 딱 맞아 떨어져야 생성됩니다 

 

• 위 코드를 실행 시키면 이렇게 이미지를 생성하고 있습니다

 

• 완료 되었다면 stylegan3폴데이 house.mp4 파일이 생성되있습니다

• 재생시켜보면 다음과 같이 원하는 시드의 이미지들이 순차적으로 나오는 비디오가 재생됩니다

 

 
 
이상으로 stylgan3 훈련과정을 마치도록 하겠습니다
stylegan3를 사용하시는데 도움이 되었으면 좋겠네요 감사합니다!