import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import os
import pickle
import datetime
class Sampling(tf.keras.layers.Layer):
def __init__(self, **kwargs):
super().__init__(**kwargs)
def call(self, inputs):
mu, log_var = inputs
epsilon = tf.keras.backend.random_normal(shape=tf.keras.backend.shape(mu), mean=0., stddev=1.)
return mu + tf.keras.backend.exp(log_var / 2) * epsilon
class VAEModel(tf.keras.models.Model):
def __init__(self, encoder, decoder, r_loss_factor, **kwargs):
super().__init__(**kwargs)
self.encoder = encoder
self.decoder = decoder
self.r_loss_factor = r_loss_factor
@tf.function
def loss_fn(self, x):
z_mean, z_log_var, z = self.encoder(x)
reconstruction = self.decoder(z)
reconstruction_loss = tf.reduce_mean(
tf.square(x - reconstruction), axis=[1,2,3]
) * self.r_loss_factor
kl_loss = tf.reduce_sum(
1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var),
axis = 1
) * (-0.5)
total_loss = reconstruction_loss + kl_loss
return total_loss, reconstruction_loss, kl_loss
@tf.function
def compute_loss_and_grads(self, x):
with tf.GradientTape() as tape:
total_loss, reconstruction_loss, kl_loss = self.loss_fn(x)
grads = tape.gradient(total_loss, self.trainable_weights)
return total_loss, reconstruction_loss, kl_loss, grads
def train_step(self, data):
if isinstance(data, tuple):
data = data[0]
total_loss, reconstruction_loss, kl_loss, grads = self.compute_loss_and_grads(data)
self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
return {
"loss": tf.math.reduce_mean(total_loss),
"reconstruction_loss": tf.math.reduce_mean(reconstruction_loss),
"kl_loss": tf.math.reduce_mean(kl_loss),
}
def call(self,inputs):
_, _, z = self.encoder(inputs)
return self.decoder(z)
class VariationalAutoEncoder():
def __init__(self,
input_dim,
encoder_conv_filters,
encoder_conv_kernel_size,
encoder_conv_strides,
decoder_conv_t_filters,
decoder_conv_t_kernel_size,
decoder_conv_t_strides,
z_dim,
r_loss_factor, ### added
use_batch_norm = False,
use_dropout = False,
epoch = 0
):
self.name = 'variational_autoencoder'
self.input_dim = input_dim
self.encoder_conv_filters = encoder_conv_filters
self.encoder_conv_kernel_size = encoder_conv_kernel_size
self.encoder_conv_strides = encoder_conv_strides
self.decoder_conv_t_filters = decoder_conv_t_filters
self.decoder_conv_t_kernel_size = decoder_conv_t_kernel_size
self.decoder_conv_t_strides = decoder_conv_t_strides
self.z_dim = z_dim
self.r_loss_factor = r_loss_factor ### added
self.use_batch_norm = use_batch_norm
self.use_dropout = use_dropout
self.epoch = epoch
self.n_layers_encoder = len(encoder_conv_filters)
self.n_layers_decoder = len(decoder_conv_t_filters)
self._build()
def _build(self):
### THE ENCODER
encoder_input = tf.keras.layers.Input(shape=self.input_dim, name='encoder_input')
x = encoder_input
for i in range(self.n_layers_encoder):
x = conv_layer = tf.keras.layers.Conv2D(
filters = self.encoder_conv_filters[i],
kernel_size = self.encoder_conv_kernel_size[i],
strides = self.encoder_conv_strides[i],
padding = 'same',
name = 'encoder_conv_' + str(i)
)(x)
if self.use_batch_norm: ### The order of layers is opposite to AutoEncoder
x = tf.keras.layers.BatchNormalization()(x) ### AE: LeakyReLU -> BatchNorm
x = tf.keras.layers.LeakyReLU()(x) ### VAE: BatchNorm -> LeakyReLU
if self.use_dropout:
x = tf.keras.layers.Dropout(rate = 0.25)(x)
shape_before_flattening = tf.keras.backend.int_shape(x)[1:]
x = tf.keras.layers.Flatten()(x)
self.mu = tf.keras.layers.Dense(self.z_dim, name='mu')(x)
self.log_var = tf.keras.layers.Dense(self.z_dim, name='log_var')(x)
self.z = Sampling(name='encoder_output')([self.mu, self.log_var])
self.encoder = tf.keras.models.Model(encoder_input, [self.mu, self.log_var, self.z], name='encoder')
### THE DECODER
decoder_input = tf.keras.layers.Input(shape=(self.z_dim,), name='decoder_input')
x = decoder_input
x = tf.keras.layers.Dense(np.prod(shape_before_flattening))(x)
x = tf.keras.layers.Reshape(shape_before_flattening)(x)
for i in range(self.n_layers_decoder):
x = conv_t_layer = tf.keras.layers.Conv2DTranspose(
filters = self.decoder_conv_t_filters[i],
kernel_size = self.decoder_conv_t_kernel_size[i],
strides = self.decoder_conv_t_strides[i],
padding = 'same',
name = 'decoder_conv_t_' + str(i)
)(x)
if i < self.n_layers_decoder - 1:
if self.use_batch_norm: ### The order of layers is opposite to AutoEncoder
x = tf.keras.layers.BatchNormalization()(x) ### AE: LeakyReLU -> BatchNorm
x = tf.keras.layers.LeakyReLU()(x) ### VAE: BatchNorm -> LeakyReLU
if self.use_dropout:
x = tf.keras.layers.Dropout(rate=0.25)(x)
else:
x = tf.keras.layers.Activation('sigmoid')(x)
decoder_output = x
self.decoder = tf.keras.models.Model(decoder_input, decoder_output, name='decoder') ### added (name)
### THE FULL AUTOENCODER
self.model = VAEModel(self.encoder, self.decoder, self.r_loss_factor)
def save(self, folder):
self.save_params(os.path.join(folder, 'params.pkl'))
self.save_weights(folder)
@staticmethod
def load(folder, epoch=None): # VariationalAutoEncoder.load(folder)
params = VariationalAutoEncoder.load_params(os.path.join(folder, 'params.pkl'))
VAE = VariationalAutoEncoder(*params)
if epoch is None:
VAE.load_weights(folder)
else:
VAE.load_weights(folder, epoch-1)
VAE.epoch = epoch
return VAE
def save_params(self, filepath):
dpath, fname = os.path.split(filepath)
if dpath != '' and not os.path.exists(dpath):
os.makedirs(dpath)
with open(filepath, 'wb') as f:
pickle.dump([
self.input_dim,
self.encoder_conv_filters,
self.encoder_conv_kernel_size,
self.encoder_conv_strides,
self.decoder_conv_t_filters,
self.decoder_conv_t_kernel_size,
self.decoder_conv_t_strides,
self.z_dim,
self.r_loss_factor,
self.use_batch_norm,
self.use_dropout,
self.epoch
], f)
@staticmethod
def load_params(filepath):
with open(filepath, 'rb') as f:
params = pickle.load(f)
return params
def save_weights(self, folder, epoch=None):
if epoch is None:
self.save_model_weights(self.encoder, os.path.join(folder, f'weights/encoder-weights.h5'))
self.save_model_weights(self.decoder, os.path.join(folder, f'weights/decoder-weights.h5'))
else:
self.save_model_weights(self.encoder, os.path.join(folder, f'weights/encoder-weights_{epoch}.h5'))
self.save_model_weights(self.decoder, os.path.join(folder, f'weights/decoder-weights_{epoch}.h5'))
def save_model_weights(self, model, filepath):
dpath, fname = os.path.split(filepath)
if dpath != '' and not os.path.exists(dpath):
os.makedirs(dpath)
model.save_weights(filepath)
def load_weights(self, folder, epoch=None):
if epoch is None:
self.encoder.load_weights(os.path.join(folder, f'weights/encoder-weights.h5'))
self.decoder.load_weights(os.path.join(folder, f'weights/decoder-weights.h5'))
else:
self.encoder.load_weights(os.path.join(folder, f'weights/encoder-weights_{epoch}.h5'))
self.decoder.load_weights(os.path.join(folder, f'weights/decoder-weights_{epoch}.h5'))
def save_images(self, imgs, filepath):
z_mean, z_log_var, z = self.encoder.predict(imgs)
reconst_imgs = self.decoder.predict(z)
txts = [ f'{p[0]:.3f}, {p[1]:.3f}' for p in z ]
AutoEncoder.showImages(imgs, reconst_imgs, txts, 1.4, 1.4, 0.5, filepath)
def compile(self, learning_rate):
self.learning_rate = learning_rate
optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
self.model.compile(optimizer=optimizer) # CAUTION!!!: loss(y_true, y_pred) function is not specified.
def train_with_fit(
self,
x_train,
batch_size,
epochs,
run_folder='run/'
):
history = self.model.fit(
x_train,
x_train,
batch_size = batch_size,
shuffle=True,
initial_epoch = self.epoch,
epochs = epochs
)
if (self.epoch < epochs):
self.epoch = epochs
if run_folder != None:
self.save(run_folder)
self.save_weights(run_folder, self.epoch-1)
return history
def train_generator_with_fit(
self,
data_flow,
epochs,
run_folder='run/'
):
history = self.model.fit(
data_flow,
initial_epoch = self.epoch,
epochs = epochs
)
if (self.epoch < epochs):
self.epoch = epochs
if run_folder != None:
self.save(run_folder)
self.save_weights(run_folder, self.epoch-1)
return history
def train_tf(
self,
x_train,
batch_size = 32,
epochs = 10,
shuffle = False,
run_folder = 'run/',
optimizer = None,
save_epoch_interval = 100,
validation_data = None
):
start_time = datetime.datetime.now()
steps = x_train.shape[0] // batch_size
total_losses = []
reconstruction_losses = []
kl_losses = []
val_total_losses = []
val_reconstruction_losses = []
val_kl_losses = []
for epoch in range(self.epoch, epochs):
epoch_loss = 0
indices = tf.range(x_train.shape[0], dtype=tf.int32)
if shuffle:
indices = tf.random.shuffle(indices)
x_ = x_train[indices]
step_total_losses = []
step_reconstruction_losses = []
step_kl_losses = []
for step in range(steps):
start = batch_size * step
end = start + batch_size
total_loss, reconstruction_loss, kl_loss, grads = self.model.compute_loss_and_grads(x_[start:end])
optimizer.apply_gradients(zip(grads, self.model.trainable_weights))
step_total_losses.append(np.mean(total_loss))
step_reconstruction_losses.append(np.mean(reconstruction_loss))
step_kl_losses.append(np.mean(kl_loss))
epoch_total_loss = np.mean(step_total_losses)
epoch_reconstruction_loss = np.mean(step_reconstruction_losses)
epoch_kl_loss = np.mean(step_kl_losses)
total_losses.append(epoch_total_loss)
reconstruction_losses.append(epoch_reconstruction_loss)
kl_losses.append(epoch_kl_loss)
val_str = ''
if not validation_data is None:
x_val = validation_data
tl, rl, kl = self.model.loss_fn(x_val)
val_tl = np.mean(tl)
val_rl = np.mean(rl)
val_kl = np.mean(kl)
val_total_losses.append(val_tl)
val_reconstruction_losses.append(val_rl)
val_kl_losses.append(val_kl)
val_str = f'val loss total {val_tl:.3f} reconstruction {val_rl:.3f} kl {val_kl:.3f} '
if (epoch+1) % save_epoch_interval == 0 and run_folder != None:
self.save(run_folder)
self.save_weights(run_folder, self.epoch)
elapsed_time = datetime.datetime.now() - start_time
print(f'{epoch+1}/{epochs} {steps} loss: total {epoch_total_loss:.3f} reconstruction {epoch_reconstruction_loss:.3f} kl {epoch_kl_loss:.3f} {val_str}{elapsed_time}')
self.epoch += 1
if run_folder != None:
self.save(run_folder)
self.save_weights(run_folder, self.epoch-1)
dic = { 'loss' : total_losses, 'reconstruction_loss' : reconstruction_losses, 'kl_loss' : kl_losses }
if not validation_data is None:
dic['val_loss'] = val_total_losses
dic['val_reconstruction_loss'] = val_reconstruction_losses
dic['val_kl_loss'] = val_kl_losses
return dic
def train_tf_generator(
self,
data_flow,
epochs = 10,
run_folder = 'run/',
optimizer = None,
save_epoch_interval = 100,
validation_data_flow = None
):
start_time = datetime.datetime.now()
steps = len(data_flow)
total_losses = []
reconstruction_losses = []
kl_losses = []
val_total_losses = []
val_reconstruction_losses = []
val_kl_losses = []
for epoch in range(self.epoch, epochs):
epoch_loss = 0
step_total_losses = []
step_reconstruction_losses = []
step_kl_losses = []
for step in range(steps):
x, _ = next(data_flow)
total_loss, reconstruction_loss, kl_loss, grads = self.model.compute_loss_and_grads(x)
optimizer.apply_gradients(zip(grads, self.model.trainable_weights))
step_total_losses.append(np.mean(total_loss))
step_reconstruction_losses.append(np.mean(reconstruction_loss))
step_kl_losses.append(np.mean(kl_loss))
epoch_total_loss = np.mean(step_total_losses)
epoch_reconstruction_loss = np.mean(step_reconstruction_losses)
epoch_kl_loss = np.mean(step_kl_losses)
total_losses.append(epoch_total_loss)
reconstruction_losses.append(epoch_reconstruction_loss)
kl_losses.append(epoch_kl_loss)
val_str = ''
if not validation_data_flow is None:
step_val_tl = []
step_val_rl = []
step_val_kl = []
for i in range(len(validation_data_flow)):
x, _ = next(validation_data_flow)
tl, rl, kl = self.model.loss_fn(x)
step_val_tl.append(np.mean(tl))
step_val_rl.append(np.mean(rl))
step_val_kl.append(np.mean(kl))
val_tl = np.mean(step_val_tl)
val_rl = np.mean(step_val_rl)
val_kl = np.mean(step_val_kl)
val_total_losses.append(val_tl)
val_reconstruction_losses.append(val_rl)
val_kl_losses.append(val_kl)
val_str = f'val loss total {val_tl:.3f} reconstruction {val_rl:.3f} kl {val_kl:.3f} '
if (epoch+1) % save_epoch_interval == 0 and run_folder != None:
self.save(run_folder)
self.save_weights(run_folder, self.epoch)
elapsed_time = datetime.datetime.now() - start_time
print(f'{epoch+1}/{epochs} {steps} loss: total {epoch_total_loss:.3f} reconstruction {epoch_reconstruction_loss:.3f} kl {epoch_kl_loss:.3f} {val_str}{elapsed_time}')
self.epoch += 1
if run_folder != None:
self.save(run_folder)
self.save_weights(run_folder, self.epoch-1)
dic = { 'loss' : total_losses, 'reconstruction_loss' : reconstruction_losses, 'kl_loss' : kl_losses }
if not validation_data_flow is None:
dic['val_loss'] = val_total_losses
dic['val_reconstruction_loss'] = val_reconstruction_losses
dic['val_kl_loss'] = val_kl_losses
return dic
@staticmethod
def showImages(imgs1, imgs2, txts, w, h, vskip=0.5, filepath=None):
n = len(imgs1)
fig, ax = plt.subplots(2, n, figsize=(w * n, (2+vskip) * h))
for i in range(n):
if n == 1:
axis = ax[0]
else:
axis = ax[0][i]
img = imgs1[i].squeeze()
axis.imshow(img, cmap='gray_r')
axis.axis('off')
axis.text(0.5, -0.35, txts[i], fontsize=10, ha='center', transform=axis.transAxes)
if n == 1:
axis = ax[1]
else:
axis = ax[1][i]
img2 = imgs2[i].squeeze()
axis.imshow(img2, cmap='gray_r')
axis.axis('off')
if not filepath is None:
dpath, fname = os.path.split(filepath)
if dpath != '' and not os.path.exists(dpath):
os.makedirs(dpath)
fig.savefig(filepath, dpi=600)
plt.close()
else:
plt.show()
@staticmethod
def plot_history(vals, labels):
colors = ['red', 'blue', 'green', 'orange', 'black', 'pink']
n = len(vals)
fig, ax = plt.subplots(1, 1, figsize=(9,4))
for i in range(n):
ax.plot(vals[i], c=colors[i], label=labels[i])
ax.legend(loc='upper right')
ax.set_xlabel('epochs')
# ax[0].set_ylabel('loss')
plt.show()
