Skip to content

GitLab

Commit d2096481 authored by lucas_miranda's avatar lucas_miranda
Browse files

Implemented version of SEQ_2_SEQ VAE based on tensorflow_probability

parent b99b3ace
Show whitespace changes
Inline Side-by-side
main.ipynb
View file @ d2096481
%% Cell type:code id: tags:
``` python
%load_ext autoreload
%autoreload 2
import warnings
warnings.filterwarnings("ignore")
```
%% Cell type:code id: tags:
``` python
#from source.utils import *
from source.preprocess import *
import pickle
import matplotlib.pyplot as plt
import pandas as pd
from collections import defaultdict
from tqdm import tqdm_notebook as tqdm
```
%% Cell type:code id: tags:parameters
``` python
path = "../../Desktop/DLC_social_1/"
```
%% Cell type:markdown id: tags:
# Set up and design the project
%% Cell type:code id: tags:
``` python
with open('{}DLC_social_1_exp_conditions.pickle'.format(path), 'rb') as handle:
Treatment_dict = pickle.load(handle)
```
%% Cell type:code id: tags:
``` python
#Which angles to compute?
bp_dict = {'B_Nose':['B_Left_ear','B_Right_ear'],
'B_Left_ear':['B_Nose','B_Right_ear','B_Center','B_Left_flank'],
'B_Right_ear':['B_Nose','B_Left_ear','B_Center','B_Right_flank'],
'B_Center':['B_Left_ear','B_Right_ear','B_Left_flank','B_Right_flank','B_Tail_base'],
'B_Left_flank':['B_Left_ear','B_Center','B_Tail_base'],
'B_Right_flank':['B_Right_ear','B_Center','B_Tail_base'],
'B_Tail_base':['B_Center','B_Left_flank','B_Right_flank']}
```
%% Cell type:code id: tags:
``` python
%%time
DLC_social_1 = project(path=path,#Path where to find the required files
smooth_alpha=0.85, #Alpha value for exponentially weighted smoothing
distances=['B_Center','B_Nose','B_Left_ear','B_Right_ear','B_Left_flank',
'B_Right_flank','B_Tail_base'],
ego=False,
angles=True,
connectivity=bp_dict,
arena='circular', #Type of arena used in the experiments
arena_dims=[380], #Dimensions of the arena. Just one if it's circular
video_format='.mp4',
table_format='.h5',
exp_conditions=Treatment_dict)
```
%%%% Output: stream
CPU times: user 2.59 s, sys: 818 ms, total: 3.41 s
Wall time: 1.1 s
%% Cell type:markdown id: tags:
# Run project
%% Cell type:code id: tags:
``` python
%%time
DLC_social_1_coords = DLC_social_1.run(verbose=True)
print(DLC_social_1_coords)
type(DLC_social_1_coords)
```
%%%% Output: stream
Loading trajectories...
Smoothing trajectories...
Computing distances...
%% Cell type:markdown id: tags:
# Generate coords
%% Cell type:code id: tags:
``` python
%%time
ptest = DLC_social_1_coords.get_coords(center=True, polar=False, speed=0, length='00:10:00')
ptest._type
```
%% Cell type:code id: tags:
``` python
%%time
dtest = DLC_social_1_coords.get_distances(speed=0, length='00:10:00')
dtest._type
```
%% Cell type:code id: tags:
``` python
%%time
atest = DLC_social_1_coords.get_angles(degrees=True, speed=0, length='00:10:00')
atest._type
```
%% Cell type:markdown id: tags:
# Visualization playground
%% Cell type:code id: tags:
``` python
#ptest.plot_heatmaps(['B_Center', 'W_Center'], i=1)
```
%% Cell type:code id: tags:
``` python
#Plot animation of trajectory over time with different smoothings
#plt.plot(ptest['Day2Test13DLC']['B_Center'].iloc[:5000]['x'],
# ptest['Day2Test13DLC']['B_Center'].iloc[:5000]['y'], label='alpha=0.85')
#plt.xlabel('x')
#plt.ylabel('y')
#plt.title('Mouse Center Trajectory using different exponential smoothings')
#plt.legend()
#plt.show()
```
%% Cell type:markdown id: tags:
# Dimensionality reduction playground
%% Cell type:code id: tags:
``` python
#pca = ptest.pca(4, 1000)
```
%% Cell type:code id: tags:
``` python
#plt.scatter(*pca[0].T)
#plt.show()
```
%% Cell type:markdown id: tags:
# Preprocessing playground
%% Cell type:code id: tags:
``` python
mtest = merge_tables(DLC_social_1_coords.get_coords(center=True, polar=True, length='00:10:00'))#,
# DLC_social_1_coords.get_distances(speed=0, length='00:10:00'),
# DLC_social_1_coords.get_angles(degrees=True, speed=0, length='00:10:00'))
```
%% Cell type:code id: tags:
``` python
#pptest = mtest.preprocess(window_size=51, filter='gaussian', sigma=10, shift=20)
```
%% Cell type:code id: tags:
``` python
pttest = mtest.preprocess(window_size=11, window_step=6, filter=None, standard_scaler=True)
pttest.shape
```
%% Cell type:code id: tags:
``` python
#plt.plot(pttest[2,:,2], label='normal')
#plt.plot(pptest[2,:,2], label='gaussian')
#plt.legend()
#plt.show()
```
%% Cell type:markdown id: tags:
# Trained models playground
%% Cell type:markdown id: tags:
### Seq 2 seq Variational Auto Encoder
%% Cell type:code id: tags:
``` python
from datetime import datetime
import tensorflow.keras as k
import tensorflow as tf
```
%% Cell type:code id: tags:
``` python
NAME = 'Baseline_VAE_short_512_10=warmup_begin'
log_dir = os.path.abspath(
"logs/fit/{}_{}".format(NAME, datetime.now().strftime("%Y%m%d-%H%M%S"))
)
tensorboard_callback = k.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)
```
%% Cell type:code id: tags:
``` python
from source.models import SEQ_2_SEQ_AE, SEQ_2_SEQ_VAE, SEQ_2_SEQ_VAEP, SEQ_2_SEQ_MMVAEP
```
%% Cell type:code id: tags:
``` python
encoder, decoder, ae = SEQ_2_SEQ_AE(pttest.shape).build()
ae.build(pttest.shape)
```
%% Cell type:code id: tags:
``` python
ae.summary()
```
%% Cell type:code id: tags:
``` python
k.backend.clear_session()
encoder, generator, vae, kl_warmup_callback, mmd_warmup_callback = SEQ_2_SEQ_VAE(pttest.shape,
loss='ELBO+MMD',
kl_warmup_epochs=10,
mmd_warmup_epochs=10).build()
vae.build(pttest.shape)
#vae.build(pttest.shape)
```
%% Cell type:code id: tags:
``` python
vae.summary()
```
%% Cell type:code id: tags:
``` python
encoder, generator, vaep, kl_warmup_callback, mmd_warmup_callback = SEQ_2_SEQ_VAEP(pttest.shape,
loss='ELBO+MMD',
kl_warmup_epochs=10,
mmd_warmup_epochs=10).build()
vaep.build(pttest.shape)
```
%% Cell type:code id: tags:
``` python
vaep.summary()
```
%% Cell type:code id: tags:
``` python
encoder, generator, gmvaep, kl_warmup_callback, mmd_warmup_callback = SEQ_2_SEQ_MMVAEP(pttest.shape,
loss='ELBO+MMD',
number_of_components=2,
kl_warmup_epochs=10,
mmd_warmup_epochs=10).build()
gmvaep.build(pttest.shape)
```
%% Cell type:code id: tags:
``` python
from tensorflow.keras.utils import plot_model
plot_model(gmvaep, show_shapes=True)
```
%% Cell type:code id: tags:
``` python
?plot_model
```
%% Cell type:code id: tags:
``` python
#np.random.shuffle(pttest)
pttrain = pttest[:-15000]
pttest = pttest[-15000:]
pttrain = pttrain[:15000]
```
%% Cell type:code id: tags:
``` python
#lr_schedule = tf.keras.callbacks.LearningRateScheduler(
# lambda epoch: 1e-3 * 10**(epoch / 20))
```
%% Cell type:code id: tags:
``` python
# tf.config.experimental_run_functions_eagerly(False)
history = vae.fit(x=pttrain[:-1], y=pttrain[:-1], epochs=500, batch_size=512, verbose=1,
history = vae.fit(x=pttrain[:-1], y=pttrain[:-1], epochs=2, batch_size=512, verbose=1,
validation_data=(pttest[:-1], pttest[:-1]),
callbacks=[tensorboard_callback, kl_warmup_callback, mmd_warmup_callback])
```
%% Cell type:code id: tags:
``` python
# tf.config.experimental_run_functions_eagerly(False)
# history = vaep.fit(x=pttrain[:-1], y=[pttrain[:-1],pttrain[1:]], epochs=500, batch_size=512, verbose=1,
# validation_data=(pttest[:-1], [pttest[:-1],pttest[1:]]),
# callbacks=[tensorboard_callback, kl_warmup_callback, mmd_warmup_callback])
```
......
source/model_utils.py
View file @ d2096481
......@@ -7,20 +7,9 @@ import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfpl = tfp.layers
# Helper functions
def sampling(args, epsilon_std=1.0, number_of_components=1, categorical=None):
z_mean, z_log_sigma = args
if number_of_components == 1:
epsilon = K.random_normal(shape=K.shape(z_mean), mean=0.0, stddev=epsilon_std)
return z_mean + K.exp(z_log_sigma) * epsilon
else:
# Implement mixture of gaussians encoding and sampling
pass
def compute_kernel(x, y):
x_size = K.shape(x)[0]
y_size = K.shape(y)[0]
......@@ -120,35 +109,20 @@ class UncorrelatedFeaturesConstraint(Constraint):
return self.weightage * self.uncorrelated_feature(x)
class KLDivergenceLayer(Layer):
""" Identity transform layer that adds KL divergence
to the final model loss.
"""
def __init__(self, beta=1.0, *args, **kwargs):
class KLDivergenceLayer(tfpl.KLDivergenceAddLoss):
def __init__(self, *args, **kwargs):
self.is_placeholder = True
self.beta = beta
super(KLDivergenceLayer, self).__init__(*args, **kwargs)
def get_config(self):
config = super().get_config().copy()
config.update({"beta": self.beta})
return config
def call(self, inputs, **kwargs):
mu, log_var = inputs
KL_batch = (
-0.5
* self.beta
* K.sum(1 + log_var - K.square(mu) - K.exp(log_var), axis=-1)
def call(self, distribution_a):
kl_batch = self._regularizer(distribution_a)
self.add_loss(kl_batch, inputs=[distribution_a])
self.add_metric(
kl_batch, aggregation="mean", name="kl_divergence",
)
self.add_metric(self._regularizer._weight, aggregation="mean", name="kl_rate")
self.add_loss(K.mean(KL_batch), inputs=inputs)
self.add_metric(KL_batch, aggregation="mean", name="kl_divergence")
self.add_metric(self.beta, aggregation="mean", name="kl_rate")
return inputs
return distribution_a
class MMDiscrepancyLayer(Layer):
......@@ -156,20 +130,21 @@ class MMDiscrepancyLayer(Layer):
to the final model loss.
"""
def __init__(self, beta=1.0, *args, **kwargs):
def __init__(self, prior, beta=1.0, *args, **kwargs):
self.is_placeholder = True
self.beta = beta
self.prior = prior
super(MMDiscrepancyLayer, self).__init__(*args, **kwargs)
def get_config(self):
config = super().get_config().copy()
config.update({"beta": self.beta})
config.update({"prior": self.prior})
return config
def call(self, z, **kwargs):
true_samples = K.random_normal(K.shape(z))
true_samples = self.prior.sample(1)
mmd_batch = self.beta * compute_mmd(true_samples, z)
self.add_loss(K.mean(mmd_batch), inputs=z)
self.add_metric(mmd_batch, aggregation="mean", name="mmd")
self.add_metric(self.beta, aggregation="mean", name="mmd_rate")
......
source/models.py
View file @ d2096481
......@@ -281,9 +281,6 @@ class SEQ_2_SEQ_VAE:
encoder = BatchNormalization()(encoder)
encoder = Model_E5(encoder)
# z_mean = Dense(self.ENCODING)(encoder)
# z_log_sigma = Dense(self.ENCODING)(encoder)
encoder = Dense(
tfpl.MultivariateNormalTriL.params_size(self.ENCODING), activation=None
)(encoder)
......@@ -302,17 +299,10 @@ class SEQ_2_SEQ_VAE:
)
)
# z_mean, z_log_sigma = KLDivergenceLayer(beta=kl_beta)([z_mean, z_log_sigma])
z = tfpl.MultivariateNormalTriL(self.ENCODING)(encoder)
# z = Lambda(sampling)([z_mean, z_log_sigma])
z = tfpl.MultivariateNormalTriL(
self.ENCODING,
activity_regularizer=(
tfpl.KLDivergenceRegularizer(self.prior, weight=kl_beta)
if "ELBO" in self.loss
else None
),
)(encoder)
if "ELBO" in self.loss:
z = KLDivergenceLayer(self.prior, weight=kl_beta)(z)
mmd_warmup_callback = False
if "MMD" in self.loss:
......@@ -327,7 +317,7 @@ class SEQ_2_SEQ_VAE:
)
)
z = MMDiscrepancyLayer(beta=mmd_beta)(z)
z = MMDiscrepancyLayer(prior=self.prior, beta=mmd_beta)(z)
# Define and instantiate generator
generator = Model_D0(z)
......@@ -388,6 +378,7 @@ class SEQ_2_SEQ_VAEP:
loss="ELBO+MMD",
kl_warmup_epochs=0,
mmd_warmup_epochs=0,
prior="standard_normal",
):
self.input_shape = input_shape
self.CONV_filters = CONV_filters
......@@ -399,9 +390,16 @@ class SEQ_2_SEQ_VAEP:
self.ENCODING = ENCODING
self.learn_rate = learn_rate
self.loss = loss
self.prior = prior
self.kl_warmup = kl_warmup_epochs
self.mmd_warmup = mmd_warmup_epochs
if self.prior == "standard_normal":
self.prior = tfd.Independent(
tfd.Normal(loc=tf.zeros(self.ENCODING), scale=1),
reinterpreted_batch_ndims=1,
)
assert (
"ELBO" in self.loss or "MMD" in self.loss
), "loss must be one of ELBO, MMD or ELBO+MMD (default)"
......@@ -496,8 +494,9 @@ class SEQ_2_SEQ_VAEP:
encoder = BatchNormalization()(encoder)
encoder = Model_E5(encoder)
z_mean = Dense(self.ENCODING)(encoder)
z_log_sigma = Dense(self.ENCODING)(encoder)
encoder = Dense(
tfpl.MultivariateNormalTriL.params_size(self.ENCODING), activation=None
)(encoder)
# Define and control custom loss functions
kl_warmup_callback = False
......@@ -512,9 +511,10 @@ class SEQ_2_SEQ_VAEP:
)
)
z_mean, z_log_sigma = KLDivergenceLayer(beta=kl_beta)([z_mean, z_log_sigma])
z = tfpl.MultivariateNormalTriL(self.ENCODING)(encoder)
z = Lambda(sampling)([z_mean, z_log_sigma])
if "ELBO" in self.loss:
z = KLDivergenceLayer(self.prior, weight=kl_beta)(z)
mmd_warmup_callback = False
if "MMD" in self.loss:
......@@ -583,7 +583,7 @@ class SEQ_2_SEQ_VAEP:
)(predictor)
# end-to-end autoencoder
encoder = Model(x, z_mean, name="SEQ_2_SEQ_VEncoder")
encoder = Model(x, z, name="SEQ_2_SEQ_VEncoder")
vaep = Model(
inputs=x, outputs=[x_decoded_mean, x_predicted_mean], name="SEQ_2_SEQ_VAE"
)
......@@ -629,6 +629,7 @@ class SEQ_2_SEQ_MMVAEP:
loss="ELBO+MMD",
kl_warmup_epochs=0,
mmd_warmup_epochs=0,
prior="standard_normal",
number_of_components=1,
):
self.input_shape = input_shape
......@@ -641,13 +642,16 @@ class SEQ_2_SEQ_MMVAEP:
self.ENCODING = ENCODING
self.learn_rate = learn_rate
self.loss = loss
self.prior = prior
self.kl_warmup = kl_warmup_epochs
self.mmd_warmup = mmd_warmup_epochs
self.number_of_components = number_of_components
assert (
self.number_of_components > 0
), "The number of components must be an integer greater than zero"
if self.prior == "standard_normal":
self.prior = tfd.Independent(
tfd.Normal(loc=tf.zeros(self.ENCODING), scale=1),
reinterpreted_batch_ndims=1,
)
assert (
"ELBO" in self.loss or "MMD" in self.loss
......@@ -743,19 +747,9 @@ class SEQ_2_SEQ_MMVAEP:
encoder = BatchNormalization()(encoder)
encoder = Model_E5(encoder)
# Categorical prior on mixture of Gaussians
categories = Dense(self.number_of_components, activation="softmax")
# Define mean and log_sigma as lists of vectors with an item per prior component
z_mean = []
z_log_sigma = []
for i in range(self.number_of_components):
z_mean.append(
Dense(self.ENCODING, name="{}_gaussian_mean".format(i + 1))(encoder)
)
z_log_sigma.append(
Dense(self.ENCODING, name="{}_gaussian_sigma".format(i + 1))(encoder)
)
encoder = Dense(
tfpl.MultivariateNormalTriL.params_size(self.ENCODING), activation=None
)(encoder)
# Define and control custom loss functions
kl_warmup_callback = False
......@@ -770,11 +764,10 @@ class SEQ_2_SEQ_MMVAEP:
)
)
z_mean, z_log_sigma = KLDivergenceLayer(beta=kl_beta)(
[z_mean[0], z_log_sigma[0]]
)
z = tfpl.MultivariateNormalTriL(self.ENCODING)(encoder)
z = Lambda(sampling)([z_mean, z_log_sigma])
if "ELBO" in self.loss:
z = KLDivergenceLayer(self.prior, weight=kl_beta)(z)
mmd_warmup_callback = False
if "MMD" in self.loss:
......@@ -803,7 +796,7 @@ class SEQ_2_SEQ_MMVAEP:
generator = Model_D5(generator)
generator = Model_B5(generator)
x_decoded_mean = TimeDistributed(
Dense(self.input_shape[2]), name="gmvaep_reconstruction"
Dense(self.input_shape[2]), name="vaep_reconstruction"
)(generator)
# Define and instantiate predictor
......@@ -839,11 +832,11 @@ class SEQ_2_SEQ_MMVAEP:
)(predictor)
predictor = BatchNormalization()(predictor)
x_predicted_mean = TimeDistributed(
Dense(self.input_shape[2]), name="gmvaep_prediction"
Dense(self.input_shape[2]), name="vaep_prediction"
)(predictor)
# end-to-end autoencoder
encoder = Model(x, z_mean, name="SEQ_2_SEQ_VEncoder")
encoder = Model(x, z, name="SEQ_2_SEQ_VEncoder")
gmvaep = Model(
inputs=x, outputs=[x_decoded_mean, x_predicted_mean], name="SEQ_2_SEQ_VAE"
)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment