def elbo_components(self, inputs, training=None, mask=None, **kwargs):
llk, kl = super().elbo_components(inputs, mask=mask, training=training)
px_z, qz_x = self.last_outputs
for z, qz in zip(as_tuple(self.latents), as_tuple(qz_x)):
tc = total_correlation(tf.convert_to_tensor(qz), qz)
kl[f'tc_{z.name}'] = (self.beta - 1.) * tc
return llk, kl
def _elbo(self, X, pX_Z, qZ_X, analytic, reverse, sample_shape, mask,
training):
llk, div = super()._elbo(X, pX_Z, qZ_X, analytic, reverse,
sample_shape)
for name, q in zip(self.latent_names, qZ_X):
tc = total_correlation(tf.convert_to_tensor(q), q)
div['tc_%s' % name] = (self.beta - 1.) * tc
return llk, div
def cal_total_correlation(self):
r""" Estimation of total correlation based on fitted Gaussian
Return:
tuple of 2 scalars: total correlation estimation for train and test set
"""
samples = [qz.sample(seed=self.randint) for qz in self.representations]
return tuple([
losses.total_correlation(z, qz).numpy()
for z, qz in zip(samples, self.representations)
])
def cal_total_correlation(self):
r""" Estimation of total correlation based on fitted Gaussian
Return:
tuple of 2 scalars: total correlation estimation for train and test set
"""
# memory complexity O(n*n*d), better do it on CPU
with tf.device('/CPU:0'):
return dict(tc=np.mean([
losses.total_correlation(qz.sample(seed=self.randint), qz).numpy()
for qz in self.representations
]))
def _elbo(
self,
inputs: Union[TensorTypes, List[TensorTypes]],
pX_Z: Union[Distribution, List[Distribution]],
qZ_X: Union[Distribution, List[Distribution]],
mask: Optional[TensorTypes] = None,
training: Optional[bool] = None
) -> Tuple[Dict[str, Tensor], Dict[str, Tensor]]:
llk, div = super()._elbo(inputs, pX_Z, qZ_X, mask=mask, training=training)
for name, q in zip(self.latent_names, qZ_X):
tc = total_correlation(q.sample(), q)
div[f'tc_{name}'] = (self.beta - 1.) * tc
return llk, div
def elbo_components(self,
inputs,
training=None,
pX_Z=None,
qZ_X=None,
mask=None):
llk, kl = super().elbo_components(inputs,
pX_Z=pX_Z,
qZ_X=qZ_X,
mask=mask,
training=training)
for z, qz in zip(self.latents, tf.nest.flatten(qZ_X)):
tc = total_correlation(tf.convert_to_tensor(qz), qz)
kl[f'tc_{z.name}'] = (self.beta - 1.) * tc
return llk, kl