Add a new model
We encourage the authors to implement their own multimodal VAE models into our toolkit. Here we describe how to do it.
General requirements
The toolkit is written in PyTorch using the PyTorch Lightning framework and we expect new models to use this framework as well. Currently, it is possible to implement unimodal VAEs and any multimodal VAEs which use dedicated VAE instances for each modality. You can add a new objective, encoder/decoder networks and of course other support modules that are needed.
Below we show a step-by-step tutorial on how to add a new model.
Adding a new model
First, we start by defining the model in mmvae_models.py. Our model will need a name and should inherit the TorchMMVAE
class defined in mmvae_base.py.
self.modelName will be used for model selection from the config file.
class POE(TorchMMVAE):
def __init__(self, vaes:list, n_latents:int, obj_config:dict, model_config=None):
"""
Multimodal Variaional Autoencoder with Product of Experts https://github.com/mhw32/multimodal-vae-public
:param vaes: list of modality-specific vae objects
:type vaes: list
:param n_latents: dimensionality of the (shared) latent space
:type n_latents: int
:param obj_cofig: config with objective-specific parameters (obj name, beta.)
:type obj_config: dict
:param model_cofig: config with model-specific parameters
:type model_config: dict
"""
super().__init__(n_latents, **obj_config)
self.vaes = nn.ModuleDict(vaes)
self.model_config = model_config
self.modelName = 'poe'
self.pz = dist.Normal
self.prior_dist = dist.Normal
The TorchMMVAE class includes the bare functional minimum for a multimodal VAE, i.e. the forward pass, encode and decode functions and modality_mixing function.
The newly added model can override these methods or keep them as they are and only add the modality_mixing method. Here we add the forward() pass and all methods necessary for the multimodal data integration. The first input parameter
will be the multimodal data specified in a config where the keys label the modalities and values contain the data (and possibly masks where applicable).
def forward(self, inputs, K=1):
"""
Forward pass that takes input data and outputs a dict with posteriors, reconstructions and latent samples
:param inputs: input data, a dict of modalities where missing modalities are replaced with None
:type inputs: dict
:param K: sample K samples from the posterior
:type K: int
:return: dict where keys are modalities and values are a named tuple
:rtype: dict
"""
mu, logvar, single_params = self.modality_mixing(inputs, K)
qz_x = dist.Normal(*[mu, logvar])
z = qz_x.rsample(torch.Size([1]))
qz_d, px_d, z_d = {}, {}, {}
for mod, vae in self.vaes.items():
px_d[mod] = vae.px_z(*vae.dec({"latents": z, "masks": inputs[mod]["masks"]}))
for key in inputs.keys():
qz_d[key] = qz_x
z_d[key] = {"latents": z, "masks": inputs[key]["masks"]}
return self.make_output_dict(single_params, px_d, z_d, joint_dist=qz_d)
def modality_mixing(self, x, K=1):
"""
Inference module, calculates the joint posterior
:param inputs: input data, a dict of modalities where missing modalities are replaced with None
:type inputs: dict
:param K: sample K samples from the posterior
:type K: int
:return: joint posterior and individual posteriors
:rtype: tuple(torch.tensor, torch.tensor, list, list)
"""
batch_size = find_out_batch_size(x)
# initialize the universal prior expert
mu, logvar = self.prior_expert((1, batch_size, self.n_latents), use_cuda=True)
single_params = {}
for m, vae in self.vaes.items():
if x[m]["data"] is not None:
mod_mu, mod_logvar = vae.enc(x[m])
single_params[m] = dist.Normal(*[mod_mu, mod_logvar])
mu = torch.cat((mu, mod_mu.unsqueeze(0)), dim=0)
logvar = torch.cat((logvar, mod_logvar.unsqueeze(0)), dim=0)
# product of experts to combine gaussians
mu, logvar = super(POE, POE).product_of_experts(mu, logvar)
return mu, logvar, single_params
def prior_expert(self, size, use_cuda=False):
"""Universal prior expert. Here we use a spherical
Gaussian: N(0, 1).
@param size: integer
dimensionality of Gaussian
@param use_cuda: boolean [default: False]
cast CUDA on variables
"""
mu = Variable(torch.zeros(size))
logvar = Variable(torch.log(torch.ones(size)))
if use_cuda:
mu, logvar = mu.to("cuda"), logvar.to("cuda")
return mu, logvar
The forward() method must return the VAEOutput object located in output_storage.py. Proper placement of the outputs inside this object is handled automatically by TorchMMVAE, you can thus call
self.make_output_dict(encoder_dist=None, decoder_dist=None, latent_samples=None, joint_dist=None, enc_dist_private=None, dec_dist_private=None, joint_decoder_dist=None, cross_decoder_dist=None). All these arguments are optional
(depends on what your objective function will need) and must be dictionaries with modality names as keys (i.e. {“mod_1: data,, “mod_2”: data2}).
Next, we need to specify the objective() function for this model which will define the training procedure.
def objective(self, mods):
"""
Objective function for PoE
:param data: input data with modalities as keys
:type data: dict
:return obj: dictionary with the obligatory "loss" key on which the model is optimized, plus any other keys that you wish to log
:rtype obj: dict
"""
lpx_zs, klds, losses = [[] for _ in range(len(mods.keys()))], [], []
mods_inputs = subsample_input_modalities(mods)
for m, mods_input in enumerate(mods_inputs):
output = self.forward(mods_input)
output_dic = output.unpack_values()
kld = self.obj_fn.calc_kld(output_dic["joint_dist"][0], self.pz(*self.pz_params.to("cuda")))
klds.append(kld.sum(-1))
loc_lpx_z = []
for mod in output.mods.keys():
px_z = output.mods[mod].decoder_dist
self.obj_fn.set_ltype(self.vaes[mod].ltype)
lpx_z = (self.obj_fn.recon_loss_fn(px_z, mods[mod]) * self.vaes[mod].llik_scaling).sum(-1)
loc_lpx_z.append(lpx_z)
if mod == "mod_{}".format(m + 1):
lpx_zs[m].append(lpx_z)
d = {"lpx_z": torch.stack(loc_lpx_z).sum(0), "kld": kld.sum(-1), "qz_x": output_dic["encoder_dist"], "zs": output_dic["latent_samples"], "pz": self.pz, "pz_params": self.pz_params}
losses.append(self.obj_fn.calculate_loss(d)["loss"])
ind_losses = [-torch.stack(m).sum() / self.vaes["mod_{}".format(idx+1)].llik_scaling for idx, m in enumerate(lpx_zs)]
obj = {"loss": torch.stack(losses).sum(), "reconstruction_loss": ind_losses, "kld": torch.stack(klds).mean(0).sum()}
return obj
In this case, we use the subsampling strategy. We retrieve outputs from the model (line 13), calculate reconstruction losses and KL-divergences. To calculate ELBO (or any other objective),
use self.obj_fn which is an instance of MultimodalObjective in objectives.py. It contains all reconstruction loss terms and objectives like ELBO or IWAE (more to be added). Using these functions helps
unifying the code parts that should be shared among models.
The objective() function must return a dictionary which includes the “loss” key and stores a 1D torch.tensor with the computed loss. This will be passed
to the optimizer. You can also add other arbitrary keys that will be automatically logged in tensorboard.
Finally, we need to add our model to the list of all models in __init__.py located in the models directory:
from .mmvae_models import MOE as moe
from .mmvae_models import POE as poe
from .mmvae_models import MoPOE as mopoe
from .mmvae_models import DMVAE as dmvae
from .vae import VAE
__all__ = [dmvae, moe, poe, mopoe, VAE]
If we need to, we can define specific encoder and decoder networks (although we can also re-use the existing ones).
Now we should be able to train using this model. We need to create a config.yml file as follows:
batch_size: 32
epochs: 700
exp_name: poe_exp
labels: ./data/mnist_svhn/labels.pkl
lr: 1e-3
beta: 1.5
mixing: poe
n_latents: 10
obj: elbo
optimizer: adam
pre_trained: null
seed: 2
viz_freq: 20
test_split: 0.1
dataset_name: mnist_svhn
modality_1:
decoder: MNIST
encoder: MNIST
mod_type: image
recon_loss: bce
path: ./data/mnist_svhn/mnist
modality_2:
decoder: SVHN
encoder: SVHN
recon_loss: bce
mod_type: image
path: ./data/mnist_svhn/svhn
You can see that we specified “poe” as our multimodal mixing model. After configuring the experiment, we can run the training:
cd multimodal-compare
python main.py --cfg ./configs/config.yml