I tried to implement conditional variational auto encoder, using variational auto encoder at the Keras website : https://keras.io/examples/generative/vae/
I added the second input to the model but I don't know how to fit two inputs to encoder. I think custom train_step() function has some problem and should be updated according to encoder that has two inputs.
I got this error:
ValueError: Layer decoder expects 2 input(s), but it received 1 input tensors. Inputs received: [<tf.Tensor 'encoder/12_enc_sampling/add:0' shape=(8, 8) dtype=float32>]
and here is my script:
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt
from tensorflow.keras.utils import plot_model
from VAEfunctions import track_loss, sample_decoding, test, latent_space_mds
# Add Graphviz to windows PATH for plot_model
import os
# os.environ["PATH"] += os.pathsep + 'C:/Program Files/Graphviz/bin/'
print("tensorflow version:",tf.__version__)
# =====================================Set initial parameters ==============================
epoch = 3
batch_size = 8
latent_dim = 8
random_sample_points = 3 #Number of random sample points from latent space
# ==========================================================================
# for seed_value in range(7, 8): # change the seed to find the best seed_value
# Seed value
seed_value = 7 #7 #8good
# 1. Set `PYTHONHASHSEED` environment variable at a fixed value
import os
os.environ['PYTHONHASHSEED'] = str(seed_value)
# 2. Set `python` built-in pseudo-random generator at a fixed value
import random
random.seed(seed_value)
# 3. Set `numpy` pseudo-random generator at a fixed value
import numpy as np
np.random.seed(seed_value)
# 4. Set the `tensorflow` pseudo-random generator at a fixed value
import tensorflow as tf
tf.random.set_seed(seed_value)
# for later versions:
# tf.compat.v1.set_random_seed(seed_value)
# 5. Configure a new global `tensorflow` session
session_conf = tf.compat.v1.ConfigProto(intra_op_parallelism_threads=1, inter_op_parallelism_threads=1)
sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph(), config=session_conf)
tf.compat.v1.keras.backend.set_session(sess)
# ========================================Load data===============================
x_train = np.load('data_x_train_3D_all_par_3D_800_64_64_64_1.npy') # load data, #shape:(800,64,64,64,1)
x_test = np.load('data_x_test_3D_all_par_3D_200_64_64_64_1.npy')
x = np.load('data_x_3D_all_par_3D_1000_64_64_64_1.npy')
input_num = x_train.shape[0]
input_d1 = x_train.shape[1]
input_d2 = x_train.shape[2]
input_d3 = x_train.shape[3]
input_channels = x_train.shape[4]
# ================================== Load one_hot_vector (Parameters) ===================
one_hot_vec_train = np.load('one_hot_vec.npy')[:800]
print(one_hot_vec_train.shape)
# ======================View a few images=======================
print("min of x_train is:", np.min(x_train), "max of x_train is:", np.max(x_train))
print("min of x is:", np.min(x), "max of x is:", np.max(x))
plt.figure(10)
fig, axes = plt.subplots(2, 2)
# fig.tight_layout()
for ax in axes.flat:
i = 1
im = ax.imshow(x_train[100*i][31][:][:], vmin=0, vmax=1)
i = i+1
cbar_ax = fig.add_axes([0.9, 0.15, 0.03, 0.70])
fig.colorbar(im, cax=cbar_ax)
plt.suptitle('Slide[31] of four input samples')
plt.show()
""""****************************** Variational Auto Encoder Model *************************"""
# ======================================= sampling layer ===============================================
class Sampling(layers.Layer):
"""Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""
# tf.random.set_seed(0)
def call(self, inputs):
z_mean, z_log_var = inputs
batch = tf.shape(z_mean)[0]
dim = tf.shape(z_mean)[1]
epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
return z_mean + tf.exp(0.5 * z_log_var) * epsilon
# ===================== Define the VAE as a `Model` with a custom `train_step` ==============================
class CVAE(keras.Model):
# tf.random.set_seed(0)
def __init__(self, encoder, decoder, **kwargs):
super(CVAE, self).__init__(**kwargs)
self.encoder = encoder
self.decoder = decoder
self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
self.reconstruction_loss_tracker = keras.metrics.Mean(
name="reconstruction_loss"
)
self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")
self.rec_tracker = keras.metrics.Mean(name="rec") ####
@property
def metrics(self):
return [
self.total_loss_tracker,
self.reconstruction_loss_tracker,
self.kl_loss_tracker,
self.rec_tracker, ####
]
def train_step(self, data):
with tf.GradientTape() as tape:
z_mean, z_log_var, z = self.encoder(data)
reconstruction = self.decoder(z)
reconstruction_loss = tf.reduce_mean(
tf.reduce_sum(
keras.losses.mean_squared_error(data, reconstruction), axis=(1, 2, 3)
)
)
rec = tf.reduce_mean( #####
tf.reduce_mean(
keras.losses.mean_squared_error(data, reconstruction), axis=(1, 2, 3)
)
)
kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
total_loss = reconstruction_loss + kl_loss
grads = tape.gradient(total_loss, self.trainable_weights)
self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
self.total_loss_tracker.update_state(total_loss)
self.reconstruction_loss_tracker.update_state(reconstruction_loss)
self.kl_loss_tracker.update_state(kl_loss)
self.rec_tracker.update_state(rec) ####
return {
"loss": self.total_loss_tracker.result(),
"reconstruction_loss": self.reconstruction_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result(),
"rec": self.rec_tracker.result() #####
}
# ===================================== Build the encoder ====================================
input_shape = (input_d1, input_d2, input_d3, input_channels) # input_shape=(64, 64, 64, 1)
encoder_input = keras.Input(shape=input_shape, name="enc_input")
input_one_hot_vec = keras.Input(shape=(27, ), name="one_hot_vec_input") # one_hot_vect input layer shape=(27)
x = layers.Conv3D(32, 3, activation="relu", padding="same", name="1_enc_conv")(encoder_input) # delete strides=2
x = layers.MaxPooling3D(name="2_enc_MaxPool")(x)
x = layers.Dropout(0.2)(x)
x = layers.Conv3D(32, 3, activation="relu", padding="same", name="3_enc_conv")(x)
x = layers.MaxPooling3D(name="4_enc_MaxPool")(x)
x = layers.Dropout(0.2)(x)
x = layers.Conv3D(32, 3, activation="relu", padding="same", name="5_enc_conv")(x)
x = layers.MaxPooling3D(name="6_enc_MaxPool")(x)
x = layers.Dropout(0.2)(x)
x = layers.Conv3D(32, 3, activation="relu", padding="same", name="7_en_conv")(x)
enc_last_conv = x.shape
x = layers.Flatten(name="8_enc_flat")(x)
# ============Add one_hot_vector here: ==============
concat_layer_enc = layers.concatenate([x, input_one_hot_vec], axis=1) # ====== concat one_hot_vec here
x = layers.Dense(16, activation="relu", name="9_enc_dense")(concat_layer_enc)
z_mean = layers.Dense(latent_dim, name="10_enc_z_mean")(x)
z_log_var = layers.Dense(latent_dim, name="11_enc_z_log_var")(x)
z = Sampling(name="12_enc_sampling")([z_mean, z_log_var])
encoder = keras.Model([encoder_input, input_one_hot_vec], [z_mean, z_log_var, z], name="encoder")
print(encoder.summary())
plot_model(encoder, to_file='Plot_Model/CVAE/encoder.jpg', show_shapes=True)
print("encoder model created")
# ===================================== Build the decoder ====================
latent_input = keras.Input(shape=(latent_dim,), name="dec_input")
concat_layer_dec = layers.concatenate([latent_input, input_one_hot_vec], axis=1) # ======= concat one_hot_vec here
x = layers.Dense(enc_last_conv[1] * enc_last_conv[2] * enc_last_conv[3] * enc_last_conv[4], activation="relu", name="1_dec_dense")(concat_layer_dec)
x = layers.Reshape((enc_last_conv[1], enc_last_conv[2], enc_last_conv[3], enc_last_conv[4]), name="2_dec")(x)
x = layers.Conv3DTranspose(32, 3, activation="relu", padding="same", name="3_dec_conv")(x)
x = layers.UpSampling3D(name="4_dec_MaxPool")(x)
x = layers.Conv3DTranspose(32, 3, activation="relu", padding="same", name="5_dec_conv")(x)
x = layers.UpSampling3D(name="6_dec_MaxPool")(x)
x = layers.Conv3DTranspose(32, 3, activation="relu", padding="same", name="7_dec_conv")(x)
x = layers.UpSampling3D(name="8_dec_MaxPool")(x)
x = layers.Conv3DTranspose(32, 3, activation="relu", padding="same", name="9_dec_conv")(x)
decoder_outputs = layers.Conv3DTranspose(1, 3, activation="sigmoid", padding="same", name="10_dec_conv_out")(x)
decoder = keras.Model([latent_input, input_one_hot_vec], decoder_outputs, name="decoder")
print(decoder.summary())
plot_model(decoder, to_file='Plot_Model/CVAE/decoder.jpg', show_shapes=True)
# =========================================== **Train the CVAE** ==================================================
random.seed(seed_value)
cvae = CVAE(encoder, decoder)
cvae.compile(optimizer=keras.optimizers.Adam())
history = cvae.fit([x_train, one_hot_vec_train], epochs=epoch, batch_size=batch_size)
track_loss(history, epoch, batch_size, latent_dim, random_sample_points, seed_value)
random_sample_list = latent_space_mds(cvae, x_train, epoch, latent_dim, random_sample_points, seed_value)
sample_decoding(cvae, epoch, latent_dim, random_sample_points, random_sample_list, seed_value)
test(cvae, x_test, 4, epoch, latent_dim)
# end of seed loop
print("finish")
Here is the model that I created: Encoder model:
and
Here is the decoder model:
Thanks!