【勉強ノート】 VAEによる画像生成

ほぼ個人メモです。

PytorchでVAEによる画像生成を試した。
GitHubでスターが多い以下のリポジトリを使用した。
GitHub - AntixK/PyTorch-VAE: A Collection of Variational Autoencoders (VAE) in PyTorch.

以下に、手順のメモを示す。

手順

データセット(CelebA dataset)ダウンロード

CelebAは、20万を超える有名人の画像を含む大規模な顔属性データセットである。
Googleドライブからダウンロードできる。
gdownを使用して、IDを指定してダウンロードする。

pip install gdown
gdown --id 1m8-EBPgi5MRubrm6iQjafK2QMHDBMSfJ

ダウンロードしたzipを解凍

unzip celeba.zip -d Data
unzip Data/celeba/img_align_celeba.zip -d Data/celeba

yamlを編集

バニラVAEを試すため、configs/vae.yamlを編集する。
GPU IDが1になっているため、

gpus: [1]

を

gpus: [0]

に変更する。

訓練実行

python run.py -c configs/vae.yaml

実行結果の最終行は以下のように表示される。

Epoch 99: 100%|███████| 2856/2856 [01:25<00:00, 33.33it/s, loss=0.0195, v_num=1]

結果確認

結果は、logsに出力される。
loss等のグラフは、tensorboardで確認できる。

tensorboard --logdir logs/VanillaVAE/version_0

画像生成のサンプルは、
logs/VanillaVAE/version_0/Samples
にpngでエポック毎に出力されている。

100エポック学習した時点のサンプル画像
f:id:TadaoYamaoka:20220205193341p:plain

なお、サンプル画像を見ると、数エポックで学習できており、100エポックも学習する必要はなかったようだ。

画像生成

学習済みモデルを使用して、1枚ずつ画像を生成するためのスクリプトは用意されていないため、自分でスクリプトを記述する必要があった。
以下は、Jupyter Notebookでの実行例である。

%matplotlib inline
import matplotlib.pyplot as plt

import torch
ckpt = torch.load('logs/VanillaVAE/version_0/checkpoints/last.ckpt')

import yaml
config = yaml.safe_load(open('configs/vae.yaml'))

from models import *
model = vae_models[config['model_params']['name']](**config['model_params'])

from experiment import VAEXperiment
experiment = VAEXperiment(model, config['exp_params'])

experiment.load_state_dict(ckpt['state_dict'])

device = torch.device("cpu")

img = experiment.model.sample(1, device).squeeze().permute(1,2,0).detach().numpy()
plt.imshow(img)

以下のように、生成した画像が表示される。
f:id:TadaoYamaoka:20220205193700p:plain

img = experiment.model.sample(1, device).squeeze().permute(1,2,0).detach().numpy()
plt.imshow(img)

をもう一度実行すると、別の画像が生成される。
f:id:TadaoYamaoka:20220205193845p:plain

モーフィング

潜在空間でサンプリングした2点を線形補完することで、モーフィングを行う。

# モーフィング
z = torch.empty(10, experiment.model.latent_dim)
src = torch.randn(2, experiment.model.latent_dim)
for i in range(10):
    alpha = i / 9
    z[i] = (1 - alpha) * src[0] + alpha * src[1]

samples = experiment.model.decode(z)

imgs = samples.permute(0,2,3,1).detach().numpy()

plt.figure(figsize=(12,8))
for i in range(10):
    plt.subplot(1, 10, i + 1)
    plt.imshow(imgs[i])

f:id:TadaoYamaoka:20220206142135p:plain

演算

データセットに付けられたラベルを利用して、ラベルに対応する潜在空間のベクトルを算出して、潜在空間でベクトルの加算を行う。

笑顔のラベル(Smiling)に対して、ベクトルの加算を行った結果は以下の通り。

from dataset import VAEDataset
data = VAEDataset(**config["data_params"], pin_memory=len(config['trainer_params']['gpus']) != 0)
data.setup()

smiling_label = data.train_dataset.attr_names.index('Smiling')

# 潜在空間でのベクトルを算出
pos_sum = torch.zeros(experiment.model.latent_dim)
neg_sum = torch.zeros(experiment.model.latent_dim)
pos_count = 0
neg_count = 0
experiment.model.eval()
for x, target in train_dataloader:
    pos_sel = target[:,smiling_label] == 1
    neg_sel = target[:,smiling_label] == 0
    mu, log_var = experiment.model.encode(x)
    z = experiment.model.reparameterize(mu, log_var).detach()
    pos_sum += torch.sum(z[pos_sel], 0)
    neg_sum += torch.sum(z[neg_sel], 0)
    pos_count += pos_sel.sum()
    neg_count += neg_sel.sum()

pos_avr = pos_sum / pos_count
neg_avr = neg_sum / neg_count
vec = pos_avr - neg_avr

# サンプリングした点に対して潜在空間でベクトルを加算
z = torch.empty(10, experiment.model.latent_dim)
src = torch.randn(experiment.model.latent_dim)
for i in range(10):
    alpha = i - 5
    z[i] = src + alpha * vec

samples = experiment.model.decode(z)

imgs = samples.permute(0,2,3,1).detach().numpy()

plt.figure(figsize=(12,8))
for i in range(10):
    plt.subplot(1, 10, i + 1)
    plt.imshow(imgs[i])

f:id:TadaoYamaoka:20220206171308p:plain