Updated 21/Nov/2021 by Yoshihisa Nitta 
Variational Auto Encoder Training for CelebA dataset with Tensorflow 2 on Google Colab¶
Train Variational Auto Encoder on CelebA dataset. See VAE_MNIST_Train.ipynb for a description of Variational Auto Encodr.
CelebA データセットに対して Variational Auto Encoder をGoogle Colab 上の Tensorflow 2 で学習する¶
CelebA データセットに対して変分オートエンコーダを学習させる。 Variational Auto Encoder の説明は VAE_MNIST_Train.ipynb を参照すること。
In [1]:
#! pip install tensorflow==2.7.0
In [2]:
%tensorflow_version 2.x
import tensorflow as tf
print(tf.__version__)
In [3]:
! nvidia-smi
! cat /proc/cpuinfo
! cat /etc/issue
! free -h
In [4]:
from google.colab import drive
drive.mount('/content/drive')
In [5]:
! ls /content/drive
Download source file from Google Drive or nw.tsuda.ac.jp¶
Basically, gdown from Google Drive.
Download from nw.tsuda.ac.jp above only if the specifications of Google Drive change and you cannot download from Google Drive.
Google Drive または nw.tsuda.ac.jp からファイルをダウンロードする¶
基本的に Google Drive から gdown してください。
Google Drive の仕様が変わってダウンロードができない場合にのみ、nw.tsuda.ac.jp からダウンロードしてください。
In [7]:
# Download source file
nw_path = './nw'
! rm -rf {nw_path}
! mkdir -p {nw_path}
if True: # from Google Drive
url_model = 'https://drive.google.com/uc?id=1ZCihR7JkMOity4wCr66ZCp-3ZOlfwwo3'
! (cd {nw_path}; gdown {url_model})
else: # from nw.tsuda.ac.jp
URL_NW = 'https://nw.tsuda.ac.jp/lec/GoogleColab/pub'
url_model = f'{URL_NW}/models/VariationalAutoEncoder.py'
! wget -nd {url_model} -P {nw_path}
In [8]:
! cat {nw_path}/VariationalAutoEncoder.py
Preparing CelebA dataset¶
Official WWW of CelebA dataset: https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
Google Drive of CelebA dataset: https://drive.google.com/drive/folders/0B7EVK8r0v71pWEZsZE9oNnFzTm8?resourcekey=0-5BR16BdXnb8hVj6CNHKzLg
img_align_celeba.zip mirrored on my Google Drive:
https://drive.google.com/uc?id=1LFKeoI-hb96jlV0K10dO1o04iQPBoFdx
CelebA データセットを用意する¶
CelebA データセットの公式ページ: https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
CelebA データセットのGoogle Drive: https://drive.google.com/drive/folders/0B7EVK8r0v71pWEZsZE9oNnFzTm8?resourcekey=0-5BR16BdXnb8hVj6CNHKzLg
自分の Google Drive 上にミラーした img_align_celeba.zip:
https://drive.google.com/uc?id=1LFKeoI-hb96jlV0K10dO1o04iQPBoFdx
In [9]:
# Download img_align_celeba.zip from GoogleDrive
MIRRORED_URL = 'https://drive.google.com/uc?id=1LFKeoI-hb96jlV0K10dO1o04iQPBoFdx'
! gdown {MIRRORED_URL}
In [72]:
! ls -l img_align_celeba.zip
In [11]:
DATA_DIR = 'data'
DATA_SUBDIR = 'img_align_celeba'
In [12]:
! rm -rf {DATA_DIR}
! unzip -d {DATA_DIR} -q {DATA_SUBDIR}.zip
In [13]:
! ls -l {DATA_DIR}/{DATA_SUBDIR} | head
! ls {DATA_DIR}/{DATA_SUBDIR} | wc
In [14]:
# paths to all the image files.
import os
import glob
import numpy as np
all_file_paths = np.array(glob.glob(os.path.join(DATA_DIR, DATA_SUBDIR, '*.jpg')))
n_all_images = len(all_file_paths)
print(n_all_images)
In [15]:
# slect some image files.
n_to_show = 10
selected_indices = np.random.choice(range(n_all_images), n_to_show)
selected_paths = all_file_paths[selected_indices]
In [16]:
# Display some images.
%matplotlib inline
import matplotlib.pyplot as plt
fig, ax = plt.subplots(1, n_to_show, figsize=(1.4 * n_to_show, 1.4))
for i, path in enumerate(selected_paths):
img = tf.keras.preprocessing.image.load_img(path)
ax[i].imshow(img)
ax[i].axis('off')
plt.show()
In [17]:
TRAIN_DATA_DIR = 'train_data'
TEST_DATA_DIR = 'test_data'
In [18]:
import os
split = 0.05
indices = np.arange(n_all_images)
np.random.shuffle(indices)
train_indices = indices[: -int(n_all_images * split)]
test_indices = indices[-int(n_all_images * split):]
! rm -rf {TRAIN_DATA_DIR} {TEST_DATA_DIR}
dst=f'{TRAIN_DATA_DIR}/celeba'
if not os.path.exists(dst):
os.makedirs(dst)
for idx in train_indices:
path = all_file_paths[idx]
dpath, fname = os.path.split(path)
os.symlink(f'../../{path}', f'{dst}/{fname}')
dst=f'{TEST_DATA_DIR}/celeba'
if not os.path.exists(dst):
os.makedirs(dst)
for idx in test_indices:
path = all_file_paths[idx]
dpath, fname = os.path.split(path)
os.symlink(f'../../{path}', f'{dst}/{fname}')
In [19]:
INPUT_DIM = (128, 128, 3)
BATCH_SIZE = 32
In [20]:
data_gen = tf.keras.preprocessing.image.ImageDataGenerator(
rescale = 1.0/255
)
data_flow = data_gen.flow_from_directory(
TRAIN_DATA_DIR,
target_size = INPUT_DIM[:2],
batch_size = BATCH_SIZE,
shuffle=True,
class_mode = 'input'
)
val_data_flow = data_gen.flow_from_directory(
TEST_DATA_DIR,
target_size = INPUT_DIM[:2],
batch_size = BATCH_SIZE,
shuffle=True,
class_mode = 'input'
)
In [21]:
print(len(data_flow))
print(len(val_data_flow))
In [22]:
# ImageDataGenerator.next() returns the same x and y when class_mode='input'
x, y = next(data_flow)
print(x[0].shape)
%matplotlib inline
import matplotlib.pyplot as plt
n_to_show = 10
fig, ax = plt.subplots(2, n_to_show, figsize=(1.4 * n_to_show, 1.4 * 2))
for i in range(n_to_show):
ax[0][i].imshow(x[i])
ax[0][i].axis('off')
ax[1][i].imshow(y[i])
ax[1][i].axis('off')
plt.show()
In [23]:
from nw.VariationalAutoEncoder import VariationalAutoEncoder
vae = VariationalAutoEncoder(
input_dim = INPUT_DIM,
encoder_conv_filters = [ 32, 64, 64, 64 ],
encoder_conv_kernel_size = [ 3, 3, 3, 3 ],
encoder_conv_strides = [ 2, 2, 2, 2 ],
decoder_conv_t_filters = [ 64, 64, 32, 3 ],
decoder_conv_t_kernel_size = [ 3, 3, 3, 3 ],
decoder_conv_t_strides = [ 2, 2, 2, 2 ],
z_dim = 200,
use_batch_norm = True,
use_dropout = True,
r_loss_factor = 10000
)
In [24]:
vae.encoder.summary()
In [25]:
vae.decoder.summary()
In [26]:
LEARNING_RATE = 0.0005
In [27]:
MAX_EPOCHS = 4
(1) Training: Use vae.model.fit()¶
Note that the loss function is not specified at the call of vae.model.compile() function.
Since it cannot be calculated simply using y_true and y_pred, the train_step() function of the VAEModel class called from fit() is used to find loss and gradients and train them.
The self.optimizer of the VAEModel class referenced in the train_step() function is the optimizer given by the compile() function.
(1) 学習する: vae.model.fit() を使う¶
vae.model.compile() 関数の呼び出しにおいて、loss関数を指定しないことに注意が必要である。
y_true と y_pred を使って単純に計算できないので、fit() から呼び出される
VAEModel クラスの
train_step() 関数でlossとgradientsを求めて、trainingする。
train_step() 関数の中で参照される VAEModel クラスの self.optimizer は compile() 関数で与えられた optimizer である。
In [28]:
save_path1 = '/content/drive/MyDrive/ColabRun/VAE_CelebA01'
In [29]:
optimizer = tf.keras.optimizers.Adam(learning_rate = LEARNING_RATE)
vae.model.compile(optimizer=optimizer)
In [30]:
history = vae.train_generator_with_fit(
data_flow,
epochs = 3,
run_folder = save_path1
)
In [31]:
print(history.history)
In [32]:
print(history.history.keys())
In [33]:
loss1_1 = history.history['loss']
rloss1_1 = history.history['reconstruction_loss']
kloss1_1 = history.history['kl_loss']
In [34]:
# Load the saved parameters and weights.
# 保存してある学習結果をロードする。
vae_work = VariationalAutoEncoder.load(save_path1)
# Display the epoch count of the model.
# training のepoch回数を表示する。
print(vae_work.epoch)
In [35]:
# Training in addition
# 追加で training する。
vae_work.model.compile(optimizer)
history2 = vae_work.train_generator_with_fit(
data_flow,
epochs = MAX_EPOCHS,
run_folder = save_path1
)
In [36]:
print(len(history2.history))
In [37]:
loss1_2 = history2.history['loss']
rloss1_2 = history2.history['reconstruction_loss']
kloss1_2 = history2.history['kl_loss']
loss1 = np.concatenate([loss1_1, loss1_2], axis=0)
rloss1 = np.concatenate([rloss1_1, rloss1_2], axis=0)
kloss1 = np.concatenate([kloss1_1, kloss1_2], axis=0)
In [38]:
VariationalAutoEncoder.plot_history([loss1, rloss1, kloss1], ['total_loss', 'reconstruct_loss', 'kl_loss'])
In [39]:
x_, _ = next(val_data_flow)
selected_images = x_[:10]
In [40]:
z_mean, z_log_var, z = vae_work.encoder(selected_images)
reconst_images = vae_work.decoder(z).numpy() # Convert Tensor to numpy array.
txts = [f'{p[0]:.3f}, {p[1]:.3f}' for p in z ]
In [41]:
%matplotlib inline
VariationalAutoEncoder.showImages(selected_images, reconst_images, txts, 1.4, 1.4)
(2) Training with tf.GradientTape() function.¶
Instead of using fit(), calculate the loss in your own train() function, find the gradients, and apply them to the variables.
The train_tf() function is speeding up by declaring <code>@tf.function</code> the compute_loss_and_grads() function.
(2) tf.GradientTape() 関数を使った学習¶
fit() 関数を使わずに、自分で記述した train() 関数内で loss を計算し、gradients を求めて、変数に適用する。
train_tf() 関数では、lossとgradientsの計算を行う compute_loss_and_grads() 関数を <code>@tf.function</code> 宣言することで高速化を図っている。
In [42]:
save_path2 = '/content/drive/MyDrive/ColabRun/VAE_CelebA02/'
In [43]:
from nw.VariationalAutoEncoder import VariationalAutoEncoder
vae2 = VariationalAutoEncoder(
input_dim = INPUT_DIM,
encoder_conv_filters = [ 32, 64, 64, 64 ],
encoder_conv_kernel_size = [ 3, 3, 3, 3 ],
encoder_conv_strides = [ 2, 2, 2, 2 ],
decoder_conv_t_filters = [ 64, 64, 32, 3 ],
decoder_conv_t_kernel_size = [ 3, 3, 3, 3 ],
decoder_conv_t_strides = [ 2, 2, 2, 2 ],
z_dim = 200,
use_batch_norm = True,
use_dropout = True,
r_loss_factor = 10000
)
In [44]:
optimizer2 = tf.keras.optimizers.Adam(learning_rate = LEARNING_RATE)
In [45]:
log2_1 = vae2.train_tf_generator(
data_flow,
epochs = 3,
run_folder = save_path2,
optimizer = optimizer2,
save_epoch_interval = 50
)
In [46]:
print(log2_1.keys())
loss2_1 = log2_1['loss']
rloss2_1 = log2_1['reconstruction_loss']
kloss2_1 = log2_1['kl_loss']
In [47]:
# Load the saved parameters and weights.
# 保存したパラメータと重みを読み込む
vae2_work = VariationalAutoEncoder.load(save_path2)
print(vae2_work.epoch)
In [48]:
# Train in addition
# 追加で training する。
log2_2 = vae2_work.train_tf_generator(
data_flow,
epochs = MAX_EPOCHS,
run_folder = save_path2,
optimizer = optimizer2,
save_epoch_interval=50
)
In [49]:
loss2_2 = log2_2['loss']
rloss2_2 = log2_2['reconstruction_loss']
kloss2_2 = log2_2['kl_loss']
In [50]:
loss2 = np.concatenate([loss2_1, loss2_2], axis=0)
rloss2 = np.concatenate([rloss2_1, rloss2_2], axis=0)
kloss2 = np.concatenate([kloss2_1, kloss2_2], axis=0)
In [51]:
VariationalAutoEncoder.plot_history(
[loss2],
['total_loss']
)
In [52]:
VariationalAutoEncoder.plot_history(
[rloss2],
['reconstruction_loss']
)
In [53]:
VariationalAutoEncoder.plot_history(
[kloss2],
['kl_loss']
)
In [54]:
z_mean2, z_log_var2, z2 = vae2_work.encoder(selected_images)
reconst_images2 = vae2_work.decoder(z2).numpy() # decoder() returns Tensor for @tf.function declaration. Convert the Tensor to numpy array.
txts2 = [f'{p[0]:.3f}, {p[1]:.3f}' for p in z2 ]
In [55]:
%matplotlib inline
VariationalAutoEncoder.showImages(selected_images, reconst_images2, txts2, 1.4, 1.4)
(3) Trainig with tf.GradientTape() function and Learning rate decay¶
Calculate the loss and gradients with the tf.GradientTape() function, and apply the gradients to the variables.
In addition, perform Learning rate decay in the optimizer.
(3) tf.GradientTape() 関数と学習率減数を使った学習¶
tf.GradientTape() 関数を使って loss と gradients を計算して、gradients を変数に適用する。
さらに、optimizer において Learning rate decay を行う。
In [56]:
save_path3 = '/content/drive/MyDrive/ColabRun/VAE_CelebA03/'
In [57]:
from nw.VariationalAutoEncoder import VariationalAutoEncoder
vae3 = VariationalAutoEncoder(
input_dim = INPUT_DIM,
encoder_conv_filters = [ 32, 64, 64, 64 ],
encoder_conv_kernel_size = [ 3, 3, 3, 3 ],
encoder_conv_strides = [ 2, 2, 2, 2 ],
decoder_conv_t_filters = [ 64, 64, 32, 3 ],
decoder_conv_t_kernel_size = [ 3, 3, 3, 3 ],
decoder_conv_t_strides = [ 2, 2, 2, 2 ],
z_dim = 200,
use_batch_norm = True,
use_dropout = True,
r_loss_factor = 10000
)
In [58]:
# initial_learning_rate * decay_rate ^ (step // decay_steps)
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate = LEARNING_RATE,
decay_steps = len(data_flow),
decay_rate=0.96
)
optimizer3 = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
In [59]:
log3_1 = vae3.train_tf_generator(
data_flow,
epochs = 3,
run_folder = save_path3,
optimizer = optimizer3,
save_epoch_interval=50,
validation_data_flow = val_data_flow
)
In [60]:
print(log3_1.keys())
loss3_1 = log3_1['loss']
rloss3_1 = log3_1['reconstruction_loss']
kloss3_1 = log3_1['kl_loss']
val_loss3_1 = log3_1['val_loss']
val_rloss3_1 = log3_1['val_reconstruction_loss']
val_kloss3_1 = log3_1['val_kl_loss']
In [61]:
# Load the parameters and model weights saved before
# 保存したパラメータと重みを読み込む
vae3_work = VariationalAutoEncoder.load(save_path3)
print(vae3_work.epoch)
In [62]:
# Training in addition
# 追加で training する。
log3_2 = vae3_work.train_tf_generator(
data_flow,
epochs = MAX_EPOCHS,
run_folder = save_path3,
optimizer = optimizer3,
save_epoch_interval=50,
validation_data_flow = val_data_flow
)
In [63]:
loss3_2 = log3_2['loss']
rloss3_2 = log3_2['reconstruction_loss']
kloss3_2 = log3_2['kl_loss']
val_loss3_2 = log3_2['val_loss']
val_rloss3_2 = log3_2['val_reconstruction_loss']
val_kloss3_2 = log3_2['val_kl_loss']
In [64]:
loss3 = np.concatenate([loss3_1, loss3_2], axis=0)
rloss3 = np.concatenate([rloss3_1, rloss3_2], axis=0)
kloss3 = np.concatenate([kloss3_1, kloss3_2], axis=0)
val_loss3 = np.concatenate([val_loss3_1, val_loss3_2], axis=0)
val_rloss3 = np.concatenate([val_rloss3_1, val_rloss3_2], axis=0)
val_kloss3 = np.concatenate([val_kloss3_1, val_kloss3_2], axis=0)
In [65]:
VariationalAutoEncoder.plot_history(
[loss3, val_loss3],
['total_loss', 'val_total_loss']
)
In [66]:
VariationalAutoEncoder.plot_history(
[rloss3, val_rloss3],
['reconstruction_loss', 'val_reconstruction_loss']
)
In [67]:
VariationalAutoEncoder.plot_history(
[kloss3, val_kloss3],
['kl_loss', 'val_kl_loss']
)
In [68]:
z_mean3, z_log_var3, z3 = vae3_work.encoder(selected_images)
reconst_images3 = vae3_work.decoder(z3).numpy() # decoder() returns Tensor for @tf.function declaration. Convert the Tensor to numpy array.
txts3 = [f'{p[0]:.3f}, {p[1]:.3f}' for p in z3 ]
In [69]:
%matplotlib inline
VariationalAutoEncoder.showImages(selected_images, reconst_images3, txts3, 1.4, 1.4)
In [73]:
# Save loss variables for future training
# 将来の学習のために loss 変数をセーブしておく
import os
import pickle
var_path = f'{save_path3}/loss_{vae3_work.epoch-1}.pkl'
dpath, fname = os.path.split(var_path)
if dpath != '' and not os.path.exists(dpath):
os.makedirs(dpath)
with open(var_path, 'wb') as f:
pickle.dump([
loss3,
rloss3,
kloss3,
val_loss3,
val_rloss3,
val_kloss3
], f)
In [76]:
! ls -l {save_path3}
In [ ]: