Adding Tybalt Class (#98)

* add main models class

import class for tybalt functionality

* add vae utility methods class

to support tybalt class import

* add dunder init

* correct name on models.py header

tybalt/models.py

@@ -0,0 +1,171 @@
+"""
+tybalt/models.py
+2017 Gregory Way
+
+Functions enabling the construction and usage of a Tybalt model
+"""
+
+import numpy as np
+import pandas as pd
+
+import tensorflow as tf
+from keras import backend as K
+from keras import optimizers
+from keras.utils import plot_model
+from keras.layers import Input, Dense, Lambda, Activation
+from keras.layers.normalization import BatchNormalization
+from keras.models import Model, Sequential
+
+from tybalt.utils.vae_utils import VariationalLayer, WarmUpCallback
+
+
+class Tybalt():
+ """
+ Training and evaluation of a tybalt model
+ """
+ def __init__(self, original_dim, latent_dim, batch_size, epochs,
+ learning_rate, kappa, epsilon_std, beta):
+ self.original_dim = original_dim
+ self.latent_dim = latent_dim
+ self.batch_size = batch_size
+ self.epochs = epochs
+ self.learning_rate = learning_rate
+ self.kappa = kappa
+ self.epsilon_std = epsilon_std
+ self.beta = beta
+
+ def _sampling(self, args):
+ """
+ Function for reparameterization trick to make model differentiable
+ """
+ # Function with args required for Keras Lambda function
+ z_mean, z_log_var = args
+
+ # Draw epsilon of the same shape from a standard normal distribution
+ epsilon = K.random_normal(shape=tf.shape(z_mean), mean=0.,
+ stddev=self.epsilon_std)
+
+ # The latent vector is non-deterministic and differentiable
+ # in respect to z_mean and z_log_var
+ z = z_mean + K.exp(z_log_var / 2) * epsilon
+ return z
+
+ def initialize_model(self):
+ """
+ Helper function to run that builds and compiles Keras layers
+ """
+ self.build_encoder_layer()
+ self.build_decoder_layer()
+ self.compile_vae()
+
+ def build_encoder_layer(self):
+ """
+ Function to build the encoder layer connections
+ """
+ # Input place holder for RNAseq data with specific input size
+ self.rnaseq_input = Input(shape=(self.original_dim, ))
+
+ # Input layer is compressed into a mean and log variance vector of
+ # size `latent_dim`. Each layer is initialized with glorot uniform
+ # weights and each step (dense connections, batch norm, and relu
+ # activation) are funneled separately.
+ # Each vector are connected to the rnaseq input tensor
+
+ # input layer to latent mean layer
+ z_mean = Dense(self.latent_dim,
+ kernel_initializer='glorot_uniform')(self.rnaseq_input)
+ z_mean_batchnorm = BatchNormalization()(z_mean)
+ self.z_mean_encoded = Activation('relu')(z_mean_batchnorm)
+
+ # input layer to latent standard deviation layer
+ z_var = Dense(self.latent_dim,
+ kernel_initializer='glorot_uniform')(self.rnaseq_input)
+ z_var_batchnorm = BatchNormalization()(z_var)
+ self.z_var_encoded = Activation('relu')(z_var_batchnorm)
+
+ # return the encoded and randomly sampled z vector
+ # Takes two keras layers as input to the custom sampling function layer
+ self.z = Lambda(self._sampling,
+ output_shape=(self.latent_dim, ))([self.z_mean_encoded,
+ self.z_var_encoded])
+
+ def build_decoder_layer(self):
+ """
+ Function to build the decoder layer connections
+ """
+ # The decoding layer is much simpler with a single layer glorot uniform
+ # initialized and sigmoid activation
+ self.decoder_model = Sequential()
+ self.decoder_model.add(Dense(self.original_dim, activation='sigmoid',
+ input_dim=self.latent_dim))
+ self.rnaseq_reconstruct = self.decoder_model(self.z)
+
+ def compile_vae(self):
+ """
+ Creates the vae layer and compiles all layer connections
+ """
+ adam = optimizers.Adam(lr=self.learning_rate)
+ vae_layer = VariationalLayer(var_layer=self.z_var_encoded,
+ mean_layer=self.z_mean_encoded,
+ original_dim=self.original_dim,
+ beta=self.beta)(
+ [self.rnaseq_input, self.rnaseq_reconstruct])
+ self.vae = Model(self.rnaseq_input, vae_layer)
+ self.vae.compile(optimizer=adam, loss=None, loss_weights=[self.beta])
+
+ def get_summary(self):
+ self.vae.summary()
+
+ def visualize_architecture(self, output_file):
+ # Visualize the connections of the custom VAE model
+ plot_model(self.vae, to_file=output_file)
+
+ def train_vae(self, train_df, test_df):
+ self.hist = self.vae.fit(np.array(train_df),
+ shuffle=True,
+ epochs=self.epochs,
+ batch_size=self.batch_size,
+ validation_data=(np.array(test_df),
+ np.array(test_df)),
+ callbacks=[WarmUpCallback(self.beta,
+ self.kappa)])
+
+ def visualize_training(self, output_file):
+ # Visualize training performance
+ history_df = pd.DataFrame(self.hist.history)
+ ax = history_df.plot()
+ ax.set_xlabel('Epochs')
+ ax.set_ylabel('VAE Loss')
+ fig = ax.get_figure()
+ fig.savefig(output_file)
+
+ def connect_layers(self):
+ # Make connections between layers to build separate encoder and decoder
+ self.encoder = Model(self.rnaseq_input, self.z_mean_encoded)
+
+ decoder_input = Input(shape=(self.latent_dim, ))
+ _x_decoded_mean = self.decoder_model(decoder_input)
+ self.decoder = Model(decoder_input, _x_decoded_mean)
+
+ def compress(self, df):
+ # Encode rnaseq into the hidden/latent representation - and save output
+ encoded_df = self.encoder.predict_on_batch(df)
+ encoded_df = pd.DataFrame(encoded_df,
+ columns=range(1, self.latent_dim + 1),
+ index=df.index)
+ return encoded_df
+
+ def get_decoder_weights(self):
+ # build a generator that can sample from the learned distribution
+ # can generate from any sampled z vector
+ weights = []
+ for layer in self.decoder.layers:
+ weights.append(layer.get_weights())
+ return(weights)
+
+ def predict(self, df):
+ return self.decoder.predict(np.array(df))
+
+ def save_models(self, encoder_file, decoder_file):
+ self.encoder.save(encoder_file)
+ self.decoder.save(decoder_file)

tybalt/utils/vae_utils.py

@@ -0,0 +1,46 @@
+from keras import backend as K
+from keras.callbacks import Callback
+from keras.layers import Layer
+from keras import metrics
+
+
+class VariationalLayer(Layer):
+ """
+ Define a custom layer that learns and performs the training
+ """
+ def __init__(self, var_layer, mean_layer, original_dim, beta, **kwargs):
+ # https://keras.io/layers/writing-your-own-keras-layers/
+ self.is_placeholder = True
+ self.var_layer = var_layer
+ self.mean_layer = mean_layer
+ self.original_dim = original_dim
+ self.beta = beta
+ super(VariationalLayer, self).__init__(**kwargs)
+
+ def vae_loss(self, x_input, x_decoded):
+ recon_loss = self.original_dim * metrics.binary_crossentropy(x_input,
+ x_decoded)
+ kl_loss = - 0.5 * K.sum(1 + self.var_layer -
+ K.square(self.mean_layer) -
+ K.exp(self.var_layer), axis=-1)
+ return K.mean(recon_loss + (K.get_value(self.beta) * kl_loss))
+
+ def call(self, inputs):
+ x, x_decoded = inputs
+ loss = self.vae_loss(x, x_decoded)
+ self.add_loss(loss, inputs=inputs)
+ # We won't actually use the output.
+ return x
+
+
+class WarmUpCallback(Callback):
+ def __init__(self, beta, kappa):
+ self.beta = beta
+ self.kappa = kappa
+
+ def on_epoch_end(self, epoch, logs={}):
+ """
+ Behavior on each epoch
+ """
+ if K.get_value(self.beta) <= 1:
+ K.set_value(self.beta, K.get_value(self.beta) + self.kappa)