def prepare_inputs(
self, inputs: Union[TensorTypes, List[TensorTypes]],
mask: TensorTypes) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
r""" Prepare the inputs for the semi-supervised VAE, three cases arise:
- Only the unlabeled data given
- Only the labeled data given
- A mixture of both unlabeled and labeled data
The `mask` is given as indicator, 1 for labeled sample and 0 for unlabeled samples
"""
n_labels = self.n_labels
n_outputs = len(self.output_layers)
inputs = tf.nest.flatten(inputs)
batch_size = _batch_size(inputs[0])
# no labels provided:
if len(inputs) == n_outputs:
X = inputs
y = None
mask = tf.cast(tf.zeros(shape=(batch_size, 1)), tf.bool)
else:
X = inputs[:-1]
y = inputs[-1]
if mask is None:
mask = tf.cast(tf.ones(shape=(batch_size, 1)), tf.bool)
# split into unlabelled and labelled data
mask = tf.reshape(mask, (-1, ))
X_unlabelled = [
tf.boolean_mask(i, tf.logical_not(mask), axis=0) for i in X
]
X_labelled = [tf.boolean_mask(i, mask, axis=0) for i in X]
# for unlabelled data
y_unlabelled = marginalize_categorical_labels(
batch_size=_batch_size(X_unlabelled[0]),
num_classes=n_labels,
dtype=inputs[0].dtype,
)
X_unlabelled = [tf.repeat(i, n_labels, axis=0) for i in X_unlabelled]
# for labelled data
if y is not None:
y_labelled = tf.boolean_mask(y, mask, axis=0)
y = tf.concat([y_unlabelled, y_labelled], axis=0)
mask = tf.cast(
tf.concat(
[
tf.zeros(shape=(_batch_size(y_unlabelled), 1)),
tf.ones(shape=(_batch_size(y_labelled), 1))
],
axis=0,
), tf.bool)
# for only unlabelled data
else:
y = y_unlabelled
mask = tf.repeat(mask, n_labels, axis=0)
X = [
tf.concat([unlab, lab], axis=0)
for unlab, lab in zip(X_unlabelled, X_labelled)
]
return X, y, mask
def prepare_inputs(self, inputs, mask):
n_labels = self.n_labels
n_outputs = len(self.output_layers)
inputs = tf.nest.flatten(inputs)
batch_size = _batch_size(inputs[0])
# no labels provided:
if len(inputs) == n_outputs:
X = inputs
y = None
mask = tf.cast(tf.zeros(shape=(batch_size, 1)), tf.bool)
else:
X = inputs[:-1]
y = inputs[-1]
if mask is None:
mask = tf.cast(tf.ones(shape=(batch_size, 1)), tf.bool)
# split into unlabelled and labelled data
mask = tf.reshape(mask, (-1, ))
X_unlabelled = [
tf.boolean_mask(i, tf.logical_not(mask), axis=0) for i in X
]
X_labelled = [tf.boolean_mask(i, mask, axis=0) for i in X]
# for unlabelled data
y_unlabelled = marginalize_categorical_labels(
batch_size=_batch_size(X_unlabelled[0]),
num_classes=n_labels,
dtype=inputs[0].dtype,
)
X_unlabelled = [tf.repeat(i, n_labels, axis=0) for i in X_unlabelled]
# for labelled data
if y is not None:
y_labelled = tf.boolean_mask(y, mask, axis=0)
y = tf.concat([y_unlabelled, y_labelled], axis=0)
mask = tf.cast(
tf.concat(
[
tf.zeros(shape=(_batch_size(y_unlabelled), 1)),
tf.ones(shape=(_batch_size(y_labelled), 1))
],
axis=0,
), tf.bool)
# for only unlabelled data
else:
y = y_unlabelled
mask = tf.repeat(mask, n_labels, axis=0)
X = [
tf.concat([unlab, lab], axis=0)
for unlab, lab in zip(X_unlabelled, X_labelled)
]
return X, y, mask
def _prepare_elbo(self, inputs, training=None, mask=None):
X, y, mask = prepare_ssl_inputs(inputs, mask=mask, n_unsupervised_inputs=1)
X_u, X_l, y_l = split_ssl_inputs(X, y, mask)
# for simplication only 1 inputs and 1 labels are supported
X_u, X_l = X_u[0], X_l[0]
if len(y_l) > 0:
y_l = y_l[0]
else:
y_l = None
# marginalize the unsupervised data
if self.marginalize:
X_u, y_u = marginalize_categorical_labels(
X=X_u,
n_classes=self.n_classes,
dtype=self.dtype,
)
else:
y_u = None
return X_u, y_u, X_l, y_l
def elbo_components(self, inputs, training=None, mask=None):
X, y, mask = prepare_ssl_inputs(
inputs, mask=mask, n_unsupervised_inputs=self.n_observation)
X_u, X_l, y_l = split_ssl_inputs(X, y, mask)
# for simplication only 1 inputs and 1 labels are supported
X_u, X_l = X_u[0], X_l[0]
if len(y_l) > 0:
y_l = y_l[0]
else:
y_l = None
# marginalize the unsupervised data
if self.marginalize:
X_u, y_u = marginalize_categorical_labels(
X=X_u,
n_classes=self.n_classes,
dtype=self.dtype,
)
else:
y_u = None
### for unlabelled data (assumed always available)
llk_u, kl_u = super().elbo_components(inputs=[X_u, y_u],
training=training)
p_u, q_u = self.last_outputs
if self.marginalize:
qy_xu = self.classify(X_u)
else:
qy_xu = q_u[-1]
p = qy_xu.probs_parameter()
llk_u = {
k + '_u': tf.reduce_sum(p * tf.expand_dims(v, axis=-1), axis=-1)
for k, v in llk_u.items()
}
kl_u = {
k + '_u': tf.reduce_sum(tf.expand_dims(p, axis=0) *
tf.expand_dims(v, axis=-1),
axis=-1)
for k, v in kl_u.items()
}
# the entropy
entropy = -tf.reduce_sum(tf.math.multiply_no_nan(tf.math.log(p), p),
axis=-1)
llk_u['H_qy'] = entropy
### for labelled data, add the discriminative objective
if y_l is not None:
llk_l, kl_l = super().elbo_components(inputs=[X_l, y_l],
training=training)
p_l, q_l = self.last_outputs
is_ss = tf.shape(y_l)[0] > 0
llk_l = {
k + '_l': tf.cond(is_ss, lambda: v, lambda: 0.)
for k, v in llk_l.items()
}
kl_l = {
k + '_l': tf.cond(is_ss, lambda: v, lambda: 0.)
for k, v in kl_l.items()
}
qy_xl = self.classify(X_l)
if self.relaxed:
y_l = tf.clip_by_value(y_l, 1e-8, 1. - 1e-8)
llk_l['llk_classifier'] = self.alpha * tf.cond(
is_ss, lambda: qy_xl.log_prob(y_l), lambda: 0.)
else:
llk_l = {}
kl_l = {}
### merge everything
llk = {k: tf.reduce_mean(v) for k, v in dict(**llk_u, **llk_l).items()}
kl = {k: tf.reduce_mean(v) for k, v in dict(**kl_u, **kl_l).items()}
return llk, kl