Fusion Model Template: Subspace-based Fusion

This tutorial will show you how to create a subspace-based fusion model.

Note

I recommend looking at How to create your own fusion model: a general template before looking at this template, as I will skip over some of the details that are covered in that template (particularly regarding documentation and idiosyncrasies of the fusion model template).

There are two types of subspace-based fusion models in this library:

1. A model that has subspace methods trained before the main prediction model. An example of this is DAETabImgMaps. This works by training the subspace method first, then using the output of the subspace method as the input to the main prediction model.

2. A model that has subspace methods (such as an autoencoder) trained simultaneously with the main prediction model. An example of this is ConcatImgLatentTabDoubleLoss. This works by implementing a joint loss function that combines the loss of the subspace method and the loss of the main prediction model.

We will look at how to create both of these types of models.

Option 1: Simultaneously-trained subspace-based fusion model

There are two differences between this type of model and the general template in How to create your own fusion model: a general template:

• Must have the attribute self.custom_loss which is the loss function used to train the subspace method (e.g. the MSELoss for an autoencoder).

• Must output the subspace method’s output in the forward method as the second list element e.g. output = [prediction, reconstruction]

Warning

Using custom loss is currently only implemented if the second modality is the reconstructed modality, e.g. the image in tabular-image fusion, or the second tabular modality in tabular-tabular fusion.

The reconstruction shape must be the same as the input shape.

Here is a diagram of an example of a simultaneously-trained subspace-based fusion model:

And here is the code for the model:

import numpy as np
import torch
import torch.nn as nn
import lightning.pytorch as pl
from fusilli.fusionmodels.base_model import ParentFusionModel


class TemplateSubspaceFusionModel(ParentFusionModel):
    """
    Template for a subspace-based fusion model that has the subspace method trained before the main prediction model.
    """

    method_name = "Template Subspace Fusion Model"
    modality_type = "tabular_tabular"
    fusion_type = "subspace"

    def __init__(self, prediction_task, data_dims, multiclass_dimensions):
        super().__init__(prediction_task, data_dims, multiclass_dimensions)

        # nn.Module: Subspace method for the second modality
        self.subspace_method_downsample = nn.Sequential(
            nn.Linear(750, 480),
            nn.ReLU(),
            nn.Linear(480, 220),
            nn.ReLU(),
            nn.Linear(220, 88),
        )
        self.subspace_method_upsample = nn.Sequential(
            nn.Linear(88, 220),
            nn.ReLU(),
            nn.Linear(220, 480),
            nn.ReLU(),
            nn.Linear(480, 750),
        )

        # nn.Module: Prediction layers.
        # Concatenating the subspace method's output with the first tabular modality data
        self.pred_model = nn.Sequential(
            nn.Linear(88 + data_dims[0], 50),
            nn.ReLU(),
            nn.Linear(50, 25),
            nn.ReLU(),
            nn.Linear(25, 5),
        )

        self.set_final_pred_layers(input_dim=5)

        # nn.Module: Custom loss function for the reconstruction
        self.custom_loss = nn.MSELoss()

    def forward(self, x):
        """
        Forward pass of the model.

        Parameters
        ----------
        x : list
            List of modalities.

        Returns
        -------
        list
            List of outputs from the model.
        """
        tabular_1 = x[0]
        tabular_2 = x[1]

        # get the subspace method's output
        subspace_output = self.subspace_method_downsample(tabular_2)
        subspace_reconstruction = self.subspace_method_upsample(subspace_output)

        # get the prediction model's output (concatenating the subspace method's output with the tabular data)
        out_fused = self.pred_model(torch.cat([tabular_1, subspace_output]))

        prediction = self.final_prediction(out_fused)

        # returning the subspace method's output as the second list element
        return [prediction, subspace_reconstruction]

Now that we’ve got the basic structure of the model, there is one additional thing to consider:

Can the model be modified?

For the most user flexibility, the model attributes should be able to be modified (such as the subspace method layers) and the model should be able to recalculate the layers of the model if the attributes are modified.

For more information on this, see Step 3 in :ref:`how_to_contribute_a_template_other_fusion`.

---

Option 2: Pre-trained subspace-based fusion model

This section will show how to create susbapce-based fusion model which involves one or more models that have to be pre-trained.

The .py file that contains the whole fusion model must have the following three things:

1. A PyTorch Lightning module which contains the subspace model architecture, e.g. class TemplateSubspaceModel(pl.LightningModule):

2. A class with the methods load_ckpt, train, and convert_to_latent, which are used to load the pre-trained model, train a latent space, and convert data to a latent space respectively. These are called when the data for the fusion model is loaded in prepare_fusion_data().

3. The fusion model class which contains the main prediction model architecture, e.g. class TemplateSubspaceFusionModel(ParentFusionModel, nn.Module): Similar to a general fusion model, this must have the methods __init__, calc_fused_layers, and forward.

This is a diagram of an example of a pre-trained subspace-based fusion model:

Let’s go through each of these in detail.

Step 1: Create the PyTorch Lightning subspace model

Might be useful to familiarise yourself with the pytorch lightning module first.

Methods that must have specific names:

• __init__: initialising with input parameters data_dims, which is a list of the data dimensions of the input data.

• forward: the forward pass of the model. Takes x as input. Must be modifiable (see Step 3 in How to create your own fusion model: a general template) for details.

• training_step: the training step of the model. Takes batch and batch_idx as input.

• validation_step: the validation step of the model. Takes batch and batch_idx as input.

• configure_optimizers: the optimiser of the model.

Methods that can have any name:

• A method that gets the latent space of the model from the input data, e.g. encode_image for an autoencoder with an image input. In our example, this is get_latent_rep.

Here’s an example of a model with a simple 2-layer autoencoder to get the latent space of the tabular data.

class TemplateSubspaceModel(pl.LightningModule):

    def __init__(self, data_dims):
        super(TemplateSubspaceModel, self).__init__()

        self.tab_shape = data_dims[0]

        self.encoder = nn.Linear(self.tab_shape, 50)
        self.decoder = nn.Linear(50, self.tab_shape)

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)

        return x

    def training_step(self, batch, batch_idx):
        output = self(batch)

        loss = nn.MSELoss()(output, batch)

        return loss

    def validation_step(self, batch, batch_idx):
        output = self(batch)

        loss = nn.MSELoss()(output, batch)

        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)

    def get_latent_rep(self, x):
        return self.encoder(x)

Step 2: create the class with the methods load_ckpt, train, and convert_to_latent

Must have a class attribute (defined before the __init__ method) subspace_models: a list of the subspace model classes.

For our example, subspace_methods = [TemplateSubspaceModel].

Must have the following methods:

• __init__: initialising with input parameters datamodule, k, max_epochs, and train_subspace. For more detailed documentation, see concat_img_latent_tab_subspace_method.

• load_ckpt: loading the pre-trained model. Takes checkpoint_path as input.

• train: training the latent space. Takes train_dataset and val_dataset as input.

• convert_to_latent: converting the data to a latent space. Takes test_dataset as input.

Let’s create the __init__ method first.

Note

The datamodule parameter is the data module that is created in prepare_fusion_data(). This is used to get the data for the subspace method.

The input arguments that we need are datamodule, k, max_epochs, and train_subspace. These are all passed to this method during prepare_fusion_data(), so we need to make sure that we have these as input arguments.

A couple things need to happen in the __init__ method:

1. Set the datamodule attribute to the input datamodule. This is accessed during utilities relating to checkpointing.

2. The subspace model, TemplateSubspaceModel, must be initialised. This is done by calling self.subspace_models[0].

3. If train_subspace is True, then the subspace model must be trained. This means that we need to:

1. Get the appropriate checkpoint path for the subspace model. This is done by calling get_checkpoint_filenames_for_subspace_models().

2. Initialise a PyTorch Lightning trainer using init_trainer().

Here’s an example of the __init__ method:

from fusilli.utils.training_utils import get_checkpoint_filenames_for_subspace_models, init_trainer


class TemplateSubspaceMethod:
    subspace_models = [TemplateSubspaceModel]

    def __init__(self, datamodule, k=None, max_epochs=1000, train_subspace=True):
        self.datamodule = datamodule

        self.autoencoder = self.subspace_models[0](datamodule.data_dims)

        if train_subspace:
            autoencoder_ckpt_list = get_checkpoint_filenames_for_subspace_models(self, k)
            # returns a list of checkpoint paths for the subspace model (length 1 for our example)

            self.trainer = init_trainer(
                logger=None,  # no logger for the subspace models
                output_paths=self.datamodule.output_paths,  # pass in the output paths dict stored in the datamodule
                max_epochs=max_epochs,  # max_epochs is an input argument
                checkpoint_filename=autoencoder_ckpt_list[0],  # checkpoint_filename is the first element of the list
            )

    # %%
    # Now let's create the ``load_ckpt`` method. This is called when we have already trained the subspace model and we are passing new data through the model, such as in :func:`~.RealsVsPreds.from_new_data`.
    #
    # The ``state_dict`` of the model must be loaded from the checkpoint.

    # ... continuing from the previous code snippet ...

    def load_ckpt(self, checkpoint_path):
        self.autoencoder.load_state_dict(torch.load(checkpoint_path[0])["state_dict"])

Onto the train method.

The following must happen in this method:

1. The data will be input as train and validation datasets and these need to be converted to dataloaders.

2. The model will be trained and tested by calling .fit and .validate() on the trainer.

3. The latent space of the train data will be calculated by calling .get_latent_rep on the model.

4. The new train data will be returned as a list of length 2: [the predictive train features, pandas dataframe of the train labels].

Warning

Be careful not to get your train and test data mixed up! Both have to be converted to the latent space but only the train dataset should be used in .fit()

Here’s an example of the train method for our example, where the second tabular modality is being converted to a latent space to be our new second tabular modality: e.g. [tab1, tab2] -> [tab1, tab2_latent]

# imports for the train method
import pandas as pd
from torch.utils.data import DataLoader


# ... continuing from the previous code snippet ...
def train(self, train_dataset, val_dataset):
    tabular1_train_features = train_dataset[:][0]
    tabular2_train_features = train_dataset[:][1]
    train_labels = train_dataset[:][2]

    tabular1_val_features = val_dataset[:][0]
    tabular2_val_features = val_dataset[:][1]
    val_labels = val_dataset[:][2]

    # setting dataloaders for the train and validation datasets of tabular 2
    train_dataloader = DataLoader(
        tabular2_train_features,
        batch_size=16,  # customise
        shuffle=False,
    )
    val_dataloader = DataLoader(
        tabular2_val_features,
        batch_size=16,  # customise
        shuffle=False,
    )

    # training the model
    self.trainer.fit(self.autoencoder, train_dataloader, val_dataloader)

    # validating the model
    self.trainer.validate(self.autoencoder, val_dataloader)

    # setting the model to evaluation mode
    self.autoencoder.eval()

    # getting the latent space of the train data
    tabular2_train_features_latent = self.autoencoder.get_latent_rep(tabular2_train_features)

    # returning the new train data
    new_pred_features = [tabular1_train_features, tabular2_train_features_latent]
    label_dataframe = pd.DataFrame(
        train_labels, columns=["prediction_label"]
    )

    return [new_pred_features, label_dataframe]

Finally, let’s create the convert_to_latent method. This is similar to the train method, except that we don’t need to train the model, only convert the input data to the already-trained latent space.

We will return the list, like in the train method, but this time there will be an additional element in the list: the list of data dimensions [tab1_dim, tab2_dim, img_dim].

In our example’s case, the data dimensions would be [tab1_dim, tab2_latent_dim, None].

# ... continuing from the previous code snippet ...
def convert_to_latent(self, test_dataset):
    pass

Step 3: create the fusion model class

Very similar to the general fusion model template in How to create your own fusion model: a general template.

Biggest difference is that we have an additional class-level attribute subspace_method, which points to the class that we created in Step 2.

This fusion model will use the data from convert_to_latent in Step 2, not the original input data.

class TemplateSubspaceFusionModel(ParentFusionModel, nn.Module):
    method_name = "Template Subspace Fusion Model - with pre-trained subspace method"
    modality_type = "tabular_tabular"
    fusion_type = "subspace"

    # class-level attribute pointing to the subspace method class
    subspace_method = TemplateSubspaceMethod

    def __init__(self, prediction_task, data_dims, multiclass_dimensions):
        ParentFusionModel.__init__(self, prediction_task, data_dims, multiclass_dimensions)

        # nn.Module: Prediction layers concatenating the latent space with the tabular data
        self.pred_model = nn.Sequential(
            nn.Linear(data_dims[0] + data_dims[1], 50),
            nn.ReLU(),
            nn.Linear(50, 25),
            nn.ReLU(),
            nn.Linear(25, 5),
        )

        # setting the final prediction layers based on the prediction type
        self.set_final_pred_layers(input_dim=5)

    def forward(self, x):
        tabular_1 = x[0]
        tabular_2 = x[1]

        # get the prediction model's output (concatenating the latent space with the tabular data)
        out_fused = self.pred_model(torch.cat([tabular_1, tabular_2], dim=1))

        # get the final prediction
        prediction = self.final_prediction(out_fused)

        # returning the prediction as the first list element
        return [prediction, ]

As with the simultaneously-trained subspace-based fusion model, we need to think about:

Note

Can the model be modified?

For the most user flexibility, the model attributes should be able to be modified (such as the subspace method layers) and the model should be able to recalculate the layers of the model if the attributes are modified.

For more information on this, see Step 3 in How to create your own fusion model: a general template.