Method forward

model.py:288–320

(self, h, stage, h_mol)

        return hazards, survival, Y_hat

    def forward(self, h, stage, h_mol):

        # H&E embedding.
        h = self.compression_layer(h)

        # Attention MIL and first-order pooling.
        A_raw, A = self.forward_attention(h) # 1xN tiles
        h_hist = A @ h #torch.Size([1, dim_embedding]) [Sum over N(aihi,1), ..., Sum over N(aihi,dim_embedding)]
        
        # Stage learnable embedding.
        stage = self.encoding_stage_net(stage)

        # Compression h_mol.
        h_mol = self.encoding_mol_net(h_mol)

        # Attention gates on each modality.
        h_hist, stage, h_mol = self.attn_modalities(h_hist, stage, h_mol)

        # Post-compressiong H&E slide embedding.
        h_hist = self.post_compression_layer_he(h_hist)

        # Fusion.
        m = self.forward_fusion(h_hist, stage, h_mol)

        # Post-compression of multimodal embedding.
        m = self.post_compression_layer(m)

        # Survival head.
        logits  = self.classifier(m)

        hazards, survival, Y_hat = self.forward_survival(logits)

        return hazards, survival, Y_hat, A_raw, m