개요
Zookiz 캐릭터 이미지를 이용해 styleGAN 모델을 학습시켰다. 선, 색감, 표정 등 Zookiz 캐릭터의 표현방식을 모델이 (얼마나) 잘 학습할 수 있는지 살펴봤다. 구글 Colab GPU를 사용하여 효율적인 학습은 어려웠다.
전처리
# resize images
def resize_image(src_img, size=(64,64), bg_color="white"):
src_img.thumbnail(size, PIL.Image.ANTIALIAS)
new_image = PIL.Image.new("RGB", size, bg_color)
new_image.paste(src_img, (int((size[0] - src_img.size[0]) / 2), int((size[1] - src_img.size[1]) / 2)))
return new_image
test_folder = '/content/drive/MyDrive/zookiz_dataset'
test_image_files = os.listdir(test_folder)
image_arrays = []
size = (256, 256)
background_color="white"
for file_idx in range(len(test_image_files)):
img = PIL.Image.open(os.path.join(test_folder, test_image_files[file_idx]))
resized_img = np.array(resize_image(img, size, background_color))
image_arrays.append(resized_img)
초상화 styleGAN
모델
styleGAN art
stylegan-art/styleganportraits.ipynb at master · ak9250/stylegan-art
시행착오
InvalidArgumentError: 2 root error(s) found.
(0) Invalid argument: Input to reshape is a tensor with 16384 values, but the requested shape requires a multiple of 32768
[[{{node GPU0/D_loss/D/4x4/MinibatchStddev/Reshape}}]]
[[TrainD/ApplyGrads0/UpdateWeights/cond/pred_id/_2051]]
(1) Invalid argument: Input to reshape is a tensor with 16384 values, but the requested shape requires a multiple of 32768
[[{{node GPU0/D_loss/D/4x4/MinibatchStddev/Reshape}}]]
0 successful operations.
0 derived errors ignored.
During handling of the above exception, another exception occurred:
InvalidArgumentError Traceback (most recent call last)
/tensorflow-1.15.2/python3.7/tensorflow_core/python/client/session.py in _do_call(self, fn, *args)
1382 '\nsession_config.graph_options.rewrite_options.'
1383 'disable_meta_optimizer = True')
-> 1384 raise type(e)(node_def, op, message)
1385
1386 def _extend_graph(self):
InvalidArgumentError: 2 root error(s) found.
(0) Invalid argument: Input to reshape is a tensor with 16384 values, but the requested shape requires a multiple of 32768
[[node GPU0/D_loss/D/4x4/MinibatchStddev/Reshape (defined at /tensorflow-1.15.2/python3.7/tensorflow_core/python/framework/ops.py:1748) ]]
[[TrainD/ApplyGrads0/UpdateWeights/cond/pred_id/_2051]]
(1) Invalid argument: Input to reshape is a tensor with 16384 values, but the requested shape requires a multiple of 32768
[[node GPU0/D_loss/D/4x4/MinibatchStddev/Reshape (defined at /tensorflow-1.15.2/python3.7/tensorflow_core/python/framework/ops.py:1748) ]]
0 successful operations.
0 derived errors ignored.
(2) total_kimg를 잘못 설정한 것도 문제였다. 학습은 while cur_nimg < total_kimg * 1000: 아래서 반복되는데, total_kimg를 충분히 크게 설정하지 않으면 학습이 실행되지 않고 종료된다.
처음 cur_nimg는 resume_kimg * 1000로 시작해 Discriminator가 학습할 때마다 minibatch만큼 값이 추가된다. 사전학습된 모델의 resume_kimg = 11125 이므로, total_kimg가 11125보다 작게 세팅되면 학습 loop가 아예 돌지 않는다. 처음 total_kimg를 단순 데이터 개수로 오인해 작은 값을 입력한 것이 문제였다.
total_kimg 의 의미는 “Total length of the training, measured in thousands of real images”로 이미지를 1000(1k)개 단위로 계산해 학습 횟수를 설정하겠다는 것이다.
(3) 또다른 가능성은 should_stop() 함수다. 학습 loop 바로 아래 if ctx.should_stop(): break 조건문이 있다. 이 조건문에서 ctx(학습 loop를 관리하는 RunContext 클래스의 helper object)가 학습할 때마다 while문을 종료할지 체크한다. ctx.should_stop()은 실행 폴더에 abort.txt가 존재할 때 True를 반환한다.
abort.txt가 어떤 시점에 생성되는지는 모르겠지만, 세션을 초기화하지 않고 곧바로 메인 학습 코드를 다시 돌렸을 때 ctx.should_stop()가 이전 학습 과정에서 생성된 abort.txt를 인지했을 가능성이 있다.
포켓몬 styleGAN
초기 세팅
아예 다른 모델을 사용해 다시 시도해보기로 했다. 사전 학습된 모델을 정리해둔 깃허브 페이지가 있어
GitHub - justinpinkney/awesome-pretrained-stylegan: A collection of pre-trained StyleGAN models to download 참고하기 좋았다. 모델 중에서 해상도가 적당히 낮고 (512*512 미만) Zookiz 일러스트와 비슷한 특성의 이미지로 학습한 모델을 찾기로 했다. 결과적으로 256*256 포켓몬 이미지 데이터로 학습한 모델을 선택했다.
사용할 모델에 맞춰 resume_kimg를 새로 맞춰줘야 했는데, 모델 pickle 파일명에 명시되어 있지 않았다. 다음 깃허브에서 total_kimg가 7991로 설정되어 있는 것에서 유추해 resume_kimg를 7991로 설정했다.
모델에 맞게 이미지를 256*256으로 설정하고, total_kimg를 적당히 큰 숫자로 설정했다. 아래는 3 ticks까지 성공한 모습이다.
Creating the run dir root: results
Creating the run dir: results/00000-sgan-custom-1gpu
Copying files to the run dir
dnnlib: Running training.training_loop.training_loop() on localhost...
Streaming data using training.dataset.TFRecordDataset...
Instructions for updating:
Use eager execution and:
`tf.data.TFRecordDataset(path)`
Dataset shape = [3, 256, 256]
Dynamic range = [0, 255]
Label size = 0
Loading networks from "/content/drive/MyDrive/Colab Notebooks/pokemon.pkl"...
G Params OutputShape WeightShape
--- --- --- ---
latents_in - (?, 512) -
labels_in - (?, 0) -
lod - () -
dlatent_avg - (512,) -
G_mapping/latents_in - (?, 512) -
G_mapping/labels_in - (?, 0) -
G_mapping/PixelNorm - (?, 512) -
G_mapping/Dense0 262656 (?, 512) (512, 512)
G_mapping/Dense1 262656 (?, 512) (512, 512)
G_mapping/Dense2 262656 (?, 512) (512, 512)
G_mapping/Dense3 262656 (?, 512) (512, 512)
G_mapping/Dense4 262656 (?, 512) (512, 512)
G_mapping/Dense5 262656 (?, 512) (512, 512)
G_mapping/Dense6 262656 (?, 512) (512, 512)
G_mapping/Dense7 262656 (?, 512) (512, 512)
G_mapping/Broadcast - (?, 14, 512) -
G_mapping/dlatents_out - (?, 14, 512) -
Truncation - (?, 14, 512) -
G_synthesis/dlatents_in - (?, 14, 512) -
G_synthesis/4x4/Const 534528 (?, 512, 4, 4) (512,)
G_synthesis/4x4/Conv 2885632 (?, 512, 4, 4) (3, 3, 512, 512)
G_synthesis/ToRGB_lod6 1539 (?, 3, 4, 4) (1, 1, 512, 3)
G_synthesis/8x8/Conv0_up 2885632 (?, 512, 8, 8) (3, 3, 512, 512)
G_synthesis/8x8/Conv1 2885632 (?, 512, 8, 8) (3, 3, 512, 512)
G_synthesis/ToRGB_lod5 1539 (?, 3, 8, 8) (1, 1, 512, 3)
G_synthesis/Upscale2D - (?, 3, 8, 8) -
G_synthesis/Grow_lod5 - (?, 3, 8, 8) -
G_synthesis/16x16/Conv0_up 2885632 (?, 512, 16, 16) (3, 3, 512, 512)
G_synthesis/16x16/Conv1 2885632 (?, 512, 16, 16) (3, 3, 512, 512)
G_synthesis/ToRGB_lod4 1539 (?, 3, 16, 16) (1, 1, 512, 3)
G_synthesis/Upscale2D_1 - (?, 3, 16, 16) -
G_synthesis/Grow_lod4 - (?, 3, 16, 16) -
G_synthesis/32x32/Conv0_up 2885632 (?, 512, 32, 32) (3, 3, 512, 512)
G_synthesis/32x32/Conv1 2885632 (?, 512, 32, 32) (3, 3, 512, 512)
G_synthesis/ToRGB_lod3 1539 (?, 3, 32, 32) (1, 1, 512, 3)
G_synthesis/Upscale2D_2 - (?, 3, 32, 32) -
G_synthesis/Grow_lod3 - (?, 3, 32, 32) -
G_synthesis/64x64/Conv0_up 1442816 (?, 256, 64, 64) (3, 3, 512, 256)
G_synthesis/64x64/Conv1 852992 (?, 256, 64, 64) (3, 3, 256, 256)
G_synthesis/ToRGB_lod2 771 (?, 3, 64, 64) (1, 1, 256, 3)
G_synthesis/Upscale2D_3 - (?, 3, 64, 64) -
G_synthesis/Grow_lod2 - (?, 3, 64, 64) -
G_synthesis/128x128/Conv0_up 426496 (?, 128, 128, 128) (3, 3, 256, 128)
G_synthesis/128x128/Conv1 279040 (?, 128, 128, 128) (3, 3, 128, 128)
G_synthesis/ToRGB_lod1 387 (?, 3, 128, 128) (1, 1, 128, 3)
G_synthesis/Upscale2D_4 - (?, 3, 128, 128) -
G_synthesis/Grow_lod1 - (?, 3, 128, 128) -
G_synthesis/256x256/Conv0_up 139520 (?, 64, 256, 256) (3, 3, 128, 64)
G_synthesis/256x256/Conv1 102656 (?, 64, 256, 256) (3, 3, 64, 64)
G_synthesis/ToRGB_lod0 195 (?, 3, 256, 256) (1, 1, 64, 3)
G_synthesis/Upscale2D_5 - (?, 3, 256, 256) -
G_synthesis/Grow_lod0 - (?, 3, 256, 256) -
G_synthesis/images_out - (?, 3, 256, 256) -
G_synthesis/lod - () -
G_synthesis/noise0 - (1, 1, 4, 4) -
G_synthesis/noise1 - (1, 1, 4, 4) -
G_synthesis/noise2 - (1, 1, 8, 8) -
G_synthesis/noise3 - (1, 1, 8, 8) -
G_synthesis/noise4 - (1, 1, 16, 16) -
G_synthesis/noise5 - (1, 1, 16, 16) -
G_synthesis/noise6 - (1, 1, 32, 32) -
G_synthesis/noise7 - (1, 1, 32, 32) -
G_synthesis/noise8 - (1, 1, 64, 64) -
G_synthesis/noise9 - (1, 1, 64, 64) -
G_synthesis/noise10 - (1, 1, 128, 128) -
G_synthesis/noise11 - (1, 1, 128, 128) -
G_synthesis/noise12 - (1, 1, 256, 256) -
G_synthesis/noise13 - (1, 1, 256, 256) -
images_out - (?, 3, 256, 256) -
--- --- --- ---
Total 26086229
D Params OutputShape WeightShape
--- --- --- ---
images_in - (?, 3, 256, 256) -
labels_in - (?, 0) -
lod - () -
FromRGB_lod0 256 (?, 64, 256, 256) (1, 1, 3, 64)
256x256/Conv0 36928 (?, 64, 256, 256) (3, 3, 64, 64)
256x256/Conv1_down 73856 (?, 128, 128, 128) (3, 3, 64, 128)
Downscale2D - (?, 3, 128, 128) -
FromRGB_lod1 512 (?, 128, 128, 128) (1, 1, 3, 128)
Grow_lod0 - (?, 128, 128, 128) -
128x128/Conv0 147584 (?, 128, 128, 128) (3, 3, 128, 128)
128x128/Conv1_down 295168 (?, 256, 64, 64) (3, 3, 128, 256)
Downscale2D_1 - (?, 3, 64, 64) -
FromRGB_lod2 1024 (?, 256, 64, 64) (1, 1, 3, 256)
Grow_lod1 - (?, 256, 64, 64) -
64x64/Conv0 590080 (?, 256, 64, 64) (3, 3, 256, 256)
64x64/Conv1_down 1180160 (?, 512, 32, 32) (3, 3, 256, 512)
Downscale2D_2 - (?, 3, 32, 32) -
FromRGB_lod3 2048 (?, 512, 32, 32) (1, 1, 3, 512)
Grow_lod2 - (?, 512, 32, 32) -
32x32/Conv0 2359808 (?, 512, 32, 32) (3, 3, 512, 512)
32x32/Conv1_down 2359808 (?, 512, 16, 16) (3, 3, 512, 512)
Downscale2D_3 - (?, 3, 16, 16) -
FromRGB_lod4 2048 (?, 512, 16, 16) (1, 1, 3, 512)
Grow_lod3 - (?, 512, 16, 16) -
16x16/Conv0 2359808 (?, 512, 16, 16) (3, 3, 512, 512)
16x16/Conv1_down 2359808 (?, 512, 8, 8) (3, 3, 512, 512)
Downscale2D_4 - (?, 3, 8, 8) -
FromRGB_lod5 2048 (?, 512, 8, 8) (1, 1, 3, 512)
Grow_lod4 - (?, 512, 8, 8) -
8x8/Conv0 2359808 (?, 512, 8, 8) (3, 3, 512, 512)
8x8/Conv1_down 2359808 (?, 512, 4, 4) (3, 3, 512, 512)
Downscale2D_5 - (?, 3, 4, 4) -
FromRGB_lod6 2048 (?, 512, 4, 4) (1, 1, 3, 512)
Grow_lod5 - (?, 512, 4, 4) -
4x4/MinibatchStddev - (?, 513, 4, 4) -
4x4/Conv 2364416 (?, 512, 4, 4) (3, 3, 513, 512)
4x4/Dense0 4194816 (?, 512) (8192, 512)
4x4/Dense1 513 (?, 1) (512, 1)
scores_out - (?, 1) -
--- --- --- ---
Total 23052353
Setting up snapshot image grid...
Setting up run dir...
Training...
tick 1 kimg 7992.0 lod 0.00 minibatch 8 time 4m 16s sec/tick 195.4 sec/kimg 190.79 maintenance 60.4 gpumem 5.1
tick 2 kimg 7993.0 lod 0.00 minibatch 8 time 7m 37s sec/tick 187.2 sec/kimg 182.83 maintenance 13.8 gpumem 5.1
tick 3 kimg 7994.1 lod 0.00 minibatch 8 time 10m 50s sec/tick 186.8 sec/kimg 182.45 maintenance 6.2 gpumem 5.1
---------------------------------------------------------------------------
KeyboardInterrupt Traceback (most recent call last)
이후 아래와 같이 파라미터를 수정했다.
size = (256, 256)
...
tick_kimg_dict = {4: 160, 8:140, 16:120, 32:100, 64:80, 128:2, 256:1, 512:1, 1024:1}): # Resolution-specific overrides.
# tick_kimg_dict = {4: 160, 8:140, 16:120, 32:100, 64:80, 128:60, 256:40, 512:30, 1024:20}): # Resolution-specific overrides.
...
resume_run_id = "/content/drive/MyDrive/Colab Notebooks/network-snapshot-007994.pkl", # Run ID or network pkl to resume training from, None = start from scratch.
resume_kimg = 7994, # Assumed training progress at the beginning. Affects reporting and training schedule.
...
# Dataset.
desc += '-custom'; dataset = EasyDict(tfrecord_dir='smalls', resolution=256); train.mirror_augment = True
# Number of GPUs.
desc += '-1gpu'; submit_config.num_gpus = 1; sched.minibatch_base = 4; sched.minibatch_dict = {4: 128, 8: 128, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8, 512: 4}
# Default options.
train.total_kimg = 25000
sched.lod_initial_resolution = 8
sched.G_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}
sched.D_lrate_dict = EasyDict(sched.G_lrate_dict)
- size : 이미지 크기를 전처리
- tick_kimg_dict : 해상도에 따라 tick에서 처리할 kimg 개수 설정 - 해상도 128 이상일 때 값을 매우 작게 만들어서 tick이 돌아가는 속도를 빠르게 만들었음.
- resume_run_id : 위에서 3 tick만큼 더 학습한 (가장 최근) 모델 사용
- resume_kimg : 7991에서 3 tick 진행된 상태이므로 7994로 설정
- resolution=256 : custom dataset에 해상도에 대한 정보 추가
- train.total_kimg : 학습이 계속 이뤄지도록 충분히 큰 값(25000) 설정
- lod_initial_resolution : 초기 Level of detail 설정 (그런데 막상 tick 돌 때는 0.0부터 시작함. 확인 필요.)
학습
Colab 환경에서 학습을 진행하기 때문에 제약 사항이 많았다. 사용자당 GPU 사용량도 제한되어 있고, 런타임도 최대 12시간이었다. 따라서 한번 학습할 때 가능한 한 오래 학습을 시킨 후에, 세션을 초기화하고 가장 최근에 저장된 모델로 학습을 재개하기로 했다.
런타임 끊김 방지
런타임 끊김 방지를 위해 아래 코드를 개발자도구 페이지 콘솔에 추가했다.
function ClickConnect(){
console.log("1분마다 코랩 연결 끊김 방지");
document.querySelector("#top-toolbar > colab-connect-button").shadowRoot.querySelector("#connect-icon")
.click();
}
setInterval(ClickConnect, 1000 * 60);
모델 저장
본 코드는 학습을 진행하면서 tick마다 이미지와 모델의 snapshot을 저장하도록 한다. 여기에 코드를 추가해서 나의 개인 드라이브에도 이미지와 모델을 저장하도록 했다. 처음에는 모델 snapshot만 저장해서, 초반 50~60 tick의 이미지 snapshot이 없다. (매우 아쉬운 점임.)
학습을 진행하다가 중간에 network_snapshot_ticks = 2로 수정해서 tick 두 번에 모델 snapshot 하나를 저장하게 했다. 모델 파일 하나에 300MB 정도 해서, tick마다 저장하는 것이 부담스러웠기 때문이다.
(내 개인 드라이브에 공간이 없는 게 문제였다. 용량이 부족할 것 같아 학습 중간중간 개인 드라이브에서 지난 모델 snapshot들을 직접 삭제해야 했다.)
image_snapshot_ticks = 1, # How often to export image snapshots?
network_snapshot_ticks = 2, # How often to export network snapshots?
...
misc.save_image_grid(grid_fakes, os.path.join('/content/drive/MyDrive/Colab Notebooks', 'fakes%06d.png' % resume_kimg), drange=drange_net, grid_size=grid_size)
pkl_my = os.path.join('/content/drive/MyDrive/Colab Notebooks', 'network-snapshot-%06d.pkl' % (cur_nimg // 1000))
misc.save_pkl((G, D, Gs), pkl_my)
결과
총 8시간을 학습시킨 모델로 임의 이미지를 생성해보았다.
아래는 학습 중 저장한 image snapshot을 gif로 생성한 결과이다. image snapshot을 드라이브에 저장하는 코드를 학습 중간에 추가해서, 학습 초반의 snapshot이 없다.
'인공지능 > computer vision' 카테고리의 다른 글
| GAN 실험(3) - styleGAN2 + pixel2style2pixel (0) | 2022.08.10 |
|---|---|
| GAN 실험(2) - layer swapping (0) | 2022.08.10 |
| few-shot GAN ada 이해하기 (0) | 2022.08.01 |
| ReStyle 이해하기 (0) | 2022.08.01 |
| p2S2p 이해하기 (0) | 2022.08.01 |