import tensorflow as tf
import numpy as np
from tensorflow.contrib.layers import fully_connected as FC_Net
_EPSILON = 1e-8
def div(x_, y_):
return tf.div(x_, y_ + _EPSILON)
def log(x_):
return tf.log(x_ + _EPSILON)
def xavier_initialization(size):
dim_ = size[0]
xavier_stddev = 1. / tf.sqrt(dim_ / 2.)
return tf.random_normal(shape=size, stddev=xavier_stddev)
### DEFINE PREDICTOR
def predictor(x_, o_dim_, o_type_, num_layers_=1, h_dim_=100, activation_fn=tf.nn.relu, keep_prob_=1.0, w_reg_=None):
'''
INPUT
x_ : (2D-tensor) input
o_dim_ : (int) output dimension
o_type_ : (string) output type one of {'continuous', 'categorical', 'binary'}
num_layers_ : (int) # of hidden layers
activation_fn_: tf activation functions
OUTPUT
o_type_ tensor
'''
if o_type_ == 'continuous':
out_fn = None
elif o_type_ == 'categorical':
out_fn = tf.nn.softmax #for classification task
elif o_type_ == 'binary':
out_fn = tf.nn.sigmoid
else:
raise ValueError('Wrong output type. The value {}!!'.format(o_type_))
if num_layers_ == 1:
out = FC_Net(inputs=x_, num_outputs=o_dim_, activation_fn=out_fn, weights_regularizer=w_reg_, scope='out')
else: #num_layers > 1
for tmp_layer in range(num_layers_-1):
if tmp_layer == 0:
net = x_
net = FC_Net(inputs=net, num_outputs=h_dim_, activation_fn=activation_fn, weights_regularizer=w_reg_, scope='layer_'+str(tmp_layer))
net = tf.nn.dropout(net, keep_prob=keep_prob_)
out = FC_Net(inputs=net, num_outputs=o_dim_, activation_fn=out_fn, weights_regularizer=w_reg_, scope='out')
return out
### DEFINE STOCHASTIC ENCODER
def stochastic_encoder(x_, o_dim_, num_layers_=1, h_dim_=100, activation_fn=tf.nn.relu, keep_prob_=1.0, w_reg_=None):
'''
INPUT
x_ : (2D-tensor) input
o_dim_ : (int) output dimension
num_layers_ : (int) # of hidden layers
activation_fn_: tf activation functions
OUTPUT
[mu,sigma] tensor
'''
if num_layers_ == 1:
out = FC_Net(inputs=x_, num_outputs=o_dim_, activation_fn=None, weights_regularizer=w_reg_, scope='out')
else: #num_layers > 1
for tmp_layer in range(num_layers_-1):
if tmp_layer == 0:
net = x_
net = FC_Net(inputs=net, num_outputs=h_dim_, activation_fn=activation_fn, weights_regularizer=w_reg_, scope='layer_'+str(tmp_layer))
net = tf.nn.dropout(net, keep_prob=keep_prob_)
out = FC_Net(inputs=net, num_outputs=o_dim_, activation_fn=None, weights_regularizer=w_reg_, scope='out')
return out
### DEFINE SUPERVISED LOSS FUNCTION
def loss_y(y_true_, y_pred_, y_type_):
if y_type_ == 'continuous':
tmp_loss = tf.reduce_sum((y_true_ - y_pred_)**2, axis=-1)
elif y_type_ == 'categorical':
tmp_loss = - tf.reduce_sum(y_true_ * log(y_pred_), axis=-1)
elif y_type_ == 'binary':
tmp_loss = - tf.reduce_sum(y_true_ * log(y_pred_) + (1.-y_true_) * log(1.-y_pred_), axis=-1)
else:
raise ValueError('Wrong output type. The value {}!!'.format(y_type_))
return tmp_loss
### DEFINE NETWORK-RELATED FUNCTIONS
def product_of_experts(mask_, mu_set_, logvar_set_):
tmp = 1.
for m in range(len(mu_set_)):
tmp += tf.reshape(mask_[:, m], [-1,1])*div(1., tf.exp(logvar_set_[m]))
poe_var = div(1., tmp)
poe_logvar = log(poe_var)
tmp = 0.
for m in range(len(mu_set_)):
tmp += tf.reshape(mask_[:, m], [-1,1])*div(1., tf.exp(logvar_set_[m]))*mu_set_[m]
poe_mu = poe_var * tmp
return poe_mu, poe_logvar
###########################################################################
#### DEFINE PROPOSED-NETWORK
class DeepIMV_AISTATS:
'''
- Add mixture mode
- Remove common/shared parts -- go back to the previous version
- Leave the consistency loss; but make sure to set gamma = 0
'''
def __init__(self, sess, name, input_dims, network_settings):
self.sess = sess
self.name = name
# INPUT/OUTPUT DIMENSIONS
self.M = len(input_dims['x_dim_set'])
self.x_dim_set = {}
for m in range(self.M):
self.x_dim_set[m] = input_dims['x_dim_set'][m]
self.y_dim = input_dims['y_dim']
self.y_type = input_dims['y_type']
self.z_dim = input_dims['z_dim'] # z_dim is equivalent to W and Z
self.steps_per_batch = input_dims['steps_per_batch']
# PREDICTOR INFO (VIEW-SPECIFC)
self.h_dim_p1 = network_settings['h_dim_p1'] #predictor hidden nodes
self.num_layers_p1 = network_settings['num_layers_p1'] #predictor layers
# PREDICTOR INFO (MULTI_VIEW)
self.h_dim_p2 = network_settings['h_dim_p2'] #predictor hidden nodes
self.num_layers_p2 = network_settings['num_layers_p2'] #predictor layers
# ENCODER INFO
self.h_dim_e = network_settings['h_dim_e'] #encoder hidden nodes
self.num_layers_e = network_settings['num_layers_e'] #encoder layers
self.fc_activate_fn = network_settings['fc_activate_fn']
self.reg_scale = network_settings['reg_scale'] #regularization
self._build_net()
def _build_net(self):
ds = tf.contrib.distributions
# with tf.name_scope(self.name):
with tf.variable_scope(self.name):
self.mb_size = tf.placeholder(tf.int32, [], name='batch_size')
self.lr_rate = tf.placeholder(tf.float32, name='learning_rate')
self.k_prob = tf.placeholder(tf.float32, name='keep_probability')
### INPUT/OUTPUT
self.x_set = {}
for m in range(self.M):
self.x_set[m] = tf.placeholder(tf.float32, [None, self.x_dim_set[m]], 'input_{}'.format(m))
self.mask = tf.placeholder(tf.float32, [None, self.M], name='mask')
self.y = tf.placeholder(tf.float32, [None, self.y_dim], name='output')
### BALANCING COEFFICIENTS
self.alpha = tf.placeholder(tf.float32, name='coef_alpha') #Consitency Loss
self.beta = tf.placeholder(tf.float32, name='coef_beta') #Information Bottleneck
if self.reg_scale == 0:
w_reg = None
else:
w_reg = tf.contrib.layers.l1_regularizer(scale=self.reg_scale)
### PRIOR
prior_z = ds.Normal(0.0, 1.0) #PoE Prior - q(z)
prior_z_set = {}
for m in range(self.M):
prior_z_set[m] = ds.Normal(0.0, 1.0) #View-Specific Prior - q(z_{m})
### STOCHASTIC ENCODER
self.h_set = {}
self.mu_z_set = {}
self.logvar_z_set = {}
for m in range(self.M):
with tf.variable_scope('encoder{}'.format(m+1)):
self.h_set[m] = stochastic_encoder(
x_=self.x_set[m], o_dim_=2*self.z_dim,
num_layers_=self.num_layers_e, h_dim_=self.h_dim_e,
activation_fn=self.fc_activate_fn, keep_prob_=self.k_prob, w_reg_=w_reg
)
self.mu_z_set[m] = self.h_set[m][:, :self.z_dim]
self.logvar_z_set[m] = self.h_set[m][:, self.z_dim:]
self.mu_z, self.logvar_z = product_of_experts(self.mask, self.mu_z_set, self.logvar_z_set)
qz = ds.Normal(self.mu_z, tf.sqrt(tf.exp(self.logvar_z)))
self.z = qz.sample()
self.zs = qz.sample(10)
qz_set = {}
self.z_set = {}
for m in range(self.M):
qz_set[m] = ds.Normal(self.mu_z_set[m], tf.sqrt(tf.exp(self.logvar_z_set[m])))
self.z_set[m] = qz_set[m].sample()
### PREDICTOR (JOINT)
with tf.variable_scope('predictor'):
self.y_hat = predictor(
x_=self.z, o_dim_=self.y_dim, o_type_=self.y_type,
num_layers_=self.num_layers_p2, h_dim_=self.h_dim_p2,
activation_fn=self.fc_activate_fn, keep_prob_=self.k_prob, w_reg_=w_reg
)
# this will generate multiple samples of y (based on multiple samples drawn from the variational encoder.
with tf.variable_scope('predictor', reuse=True):
self.y_hats = predictor(
x_=self.zs, o_dim_=self.y_dim, o_type_=self.y_type,
num_layers_=self.num_layers_p2, h_dim_=self.h_dim_p2,
activation_fn=self.fc_activate_fn, keep_prob_=self.k_prob, w_reg_=w_reg
)
### PREDICTOR
self.y_hat_set = {}
for m in range(self.M):
with tf.variable_scope('predictor_set{}'.format(m)):
self.y_hat_set[m] = predictor(
x_=self.z_set[m], o_dim_=self.y_dim, o_type_=self.y_type,
num_layers_=self.num_layers_p1, h_dim_=self.h_dim_p1,
activation_fn=self.fc_activate_fn, keep_prob_=self.k_prob, w_reg_=w_reg
)
### OPTIMIZER
global_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
enc_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=self.name + '/encoder')
pred_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=self.name + '/predictor')
### CONSITENCY LOSS
self.LOSS_CONSISTENCY = 0.
for m in range(self.M):
self.LOSS_CONSISTENCY += 1./self.M * div(
tf.reduce_sum(self.mask[:, m] * tf.reduce_sum(ds.kl_divergence(qz, qz_set[m]), axis=-1)),
tf.reduce_sum(self.mask[:, m])
)
self.LOSS_KL = tf.reduce_mean(
tf.reduce_sum(ds.kl_divergence(qz, prior_z), axis=-1)
)
self.LOSS_P = tf.reduce_mean(loss_y(self.y, self.y_hat, self.y_type))
self.LOSS_IB_JOINT = self.LOSS_P + self.beta*self.LOSS_KL
self.LOSS_Ps_all = []
self.LOSS_KLs_all = []
for m in range(self.M):
tmp_p = loss_y(self.y, self.y_hat_set[m], self.y_type)
tmp_kl = tf.reduce_sum(ds.kl_divergence(qz_set[m], prior_z_set[m]), axis=-1)
self.LOSS_Ps_all += [div(tf.reduce_sum(self.mask[:,m]*tmp_p), tf.reduce_sum(self.mask[:,m]))]
self.LOSS_KLs_all += [div(tf.reduce_sum(self.mask[:,m]*tmp_kl), tf.reduce_sum(self.mask[:,m]))]
self.LOSS_Ps_all = tf.stack(self.LOSS_Ps_all, axis=0)
self.LOSS_KLs_all = tf.stack(self.LOSS_KLs_all, axis=0)
self.LOSS_Ps = tf.reduce_sum(self.LOSS_Ps_all)
self.LOSS_KLs = tf.reduce_sum(self.LOSS_KLs_all)
self.LOSS_IB_MARGINAL = self.LOSS_Ps + self.beta*self.LOSS_KLs
self.LOSS_TOTAL = self.LOSS_IB_JOINT\
+ self.alpha*(self.LOSS_IB_MARGINAL)\
+ tf.losses.get_regularization_loss()
self.global_step = tf.contrib.framework.get_or_create_global_step()
self.lr_rate_decayed = tf.train.exponential_decay(self.lr_rate, self.global_step,
decay_steps=2*self.steps_per_batch,
decay_rate=0.97, staircase=True)
opt = tf.train.AdamOptimizer(self.lr_rate_decayed, 0.5)
ma = tf.train.ExponentialMovingAverage(0.999, zero_debias=True)
ma_update = ma.apply(tf.model_variables())
self.solver = tf.contrib.training.create_train_op(self.LOSS_TOTAL, opt,
self.global_step,
update_ops=[ma_update])
def train(self, x_set_, y_, m_, alpha_, beta_, lr_train, k_prob=1.0):
feed_dict_ = self.make_feed_dict(x_set_)
feed_dict_.update({self.y: y_, self.mask: m_,
self.alpha: alpha_, self.beta: beta_,
self.mb_size: np.shape(x_set_[0])[0],
self.lr_rate: lr_train, self.k_prob: k_prob})
return self.sess.run([self.solver, self.LOSS_TOTAL, self.LOSS_P, self.LOSS_KL, self.LOSS_Ps,
self.LOSS_KLs, self.LOSS_CONSISTENCY],
feed_dict=feed_dict_)
def get_loss(self, x_set_, y_, m_, alpha_, beta_):
feed_dict_ = self.make_feed_dict(x_set_)
feed_dict_.update({self.y: y_, self.mask: m_,
self.alpha: alpha_, self.beta: beta_,
self.mb_size: np.shape(x_set_[0])[0], self.k_prob: 1.0})
return self.sess.run([self.LOSS_TOTAL, self.LOSS_P, self.LOSS_KL, self.LOSS_Ps,
self.LOSS_KLs, self.LOSS_CONSISTENCY, self.LOSS_Ps_all, self.LOSS_KLs_all],
feed_dict=feed_dict_)
def predict_y(self, x_set_, m_):
feed_dict_ = self.make_feed_dict(x_set_)
feed_dict_.update({self.mask: m_, self.mb_size: np.shape(x_set_[0])[0], self.k_prob: 1.0})
return self.sess.run(self.y_hat, feed_dict=feed_dict_)
def predict_ys(self, x_set_, m_):
feed_dict_ = self.make_feed_dict(x_set_)
feed_dict_.update({self.mask: m_, self.mb_size: np.shape(x_set_[0])[0], self.k_prob: 1.0})
return self.sess.run([self.y_hat, self.y_hats], feed_dict=feed_dict_)
def predict_yhat_set(self, x_set_, m_):
feed_dict_ = self.make_feed_dict(x_set_)
feed_dict_.update({self.mask: m_, self.mb_size: np.shape(x_set_[0])[0], self.k_prob: 1.0})
return self.sess.run(self.y_hat_set, feed_dict=feed_dict_)
def predict_mu_z_and_mu_z_set(self, x_set_, m_): #this outputs mu and mu_set
feed_dict_ = self.make_feed_dict(x_set_)
feed_dict_.update({self.mask: m_, self.mb_size: np.shape(x_set_[0])[0], self.k_prob: 1.0})
return self.sess.run([self.mu_z, self.mu_z_set], feed_dict=feed_dict_)
def predict_logvar_z_and_logvar_z_set(self, x_set_, m_): #this outputs sigma and sigma_set
feed_dict_ = self.make_feed_dict(x_set_)
feed_dict_.update({self.mask: m_, self.mb_size: np.shape(x_set_[0])[0], self.k_prob: 1.0})
return self.sess.run([self.logvar_z, self.logvar_z_set], feed_dict=feed_dict_)
def predict_z_n_z_set(self, x_set_, m_): #this outputs z and z_set
feed_dict_ = self.make_feed_dict(x_set_)
feed_dict_.update({self.mask: m_, self.mb_size: np.shape(x_set_[0])[0], self.k_prob: 1.0})
return self.sess.run([self.z, self.z_set], feed_dict=feed_dict_)
def make_feed_dict(self, x_set_):
feed_dict_ = {}
for m in range(len(self.x_set)):
feed_dict_[self.x_set[m]] = x_set_[m]
return feed_dict_
Datasets
Models
