from collections import OrderedDict
from os.path import join
import pdb
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from models.model_utils import *
################################
### Deep Sets Implementation ###
################################
class MIL_Sum_FC_surv(nn.Module):
def __init__(self, omic_input_dim=None, fusion=None, size_arg = "small", dropout=0.25, n_classes=4):
r"""
Deep Sets Implementation.
Args:
omic_input_dim (int): Dimension size of genomic features.
fusion (str): Fusion method (Choices: concat, bilinear, or None)
size_arg (str): Size of NN architecture (Choices: small or large)
dropout (float): Dropout rate
n_classes (int): Output shape of NN
"""
super(MIL_Sum_FC_surv, self).__init__()
self.fusion = fusion
self.size_dict_path = {"small": [1024, 512, 256], "big": [1024, 512, 384]}
self.size_dict_omic = {'small': [256, 256]}
### Deep Sets Architecture Construction
size = self.size_dict_path[size_arg]
self.phi = nn.Sequential(*[nn.Linear(size[0], size[1]), nn.ReLU(), nn.Dropout(dropout)])
self.rho = nn.Sequential(*[nn.Linear(size[1], size[2]), nn.ReLU(), nn.Dropout(dropout)])
### Constructing Genomic SNN
if self.fusion != None:
hidden = [256, 256]
fc_omic = [SNN_Block(dim1=omic_input_dim, dim2=hidden[0])]
for i, _ in enumerate(hidden[1:]):
fc_omic.append(SNN_Block(dim1=hidden[i], dim2=hidden[i+1], dropout=0.25))
self.fc_omic = nn.Sequential(*fc_omic)
if self.fusion == 'concat':
self.mm = nn.Sequential(*[nn.Linear(256*2, size[2]), nn.ReLU(), nn.Linear(size[2], size[2]), nn.ReLU()])
elif self.fusion == 'bilinear':
self.mm = BilinearFusion(dim1=256, dim2=256, scale_dim1=8, scale_dim2=8, mmhid=256)
else:
self.mm = None
self.classifier = nn.Linear(size[2], n_classes)
def relocate(self):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() >= 1:
device_ids = list(range(torch.cuda.device_count()))
self.phi = nn.DataParallel(self.phi, device_ids=device_ids).to('cuda:0')
if self.fusion is not None:
self.fc_omic = self.fc_omic.to(device)
self.mm = self.mm.to(device)
self.rho = self.rho.to(device)
self.classifier = self.classifier.to(device)
def forward(self, **kwargs):
x_path = kwargs['x_path']
h_path = self.phi(x_path).sum(axis=0)
h_path = self.rho(h_path)
if self.fusion is not None:
x_omic = kwargs['x_omic']
h_omic = self.fc_omic(x_omic).squeeze(dim=0)
if self.fusion == 'bilinear':
h = self.mm(h_path.unsqueeze(dim=0), h_omic.unsqueeze(dim=0)).squeeze()
elif self.fusion == 'concat':
h = self.mm(torch.cat([h_path, h_omic], axis=0))
else:
h = h_path # [256] vector
logits = self.classifier(h).unsqueeze(0) # logits needs to be a [1 x 4] vector
Y_hat = torch.topk(logits, 1, dim = 1)[1]
hazards = torch.sigmoid(logits)
S = torch.cumprod(1 - hazards, dim=1)
return hazards, S, Y_hat, None, None
################################
# Attention MIL Implementation #
################################
class MIL_Attention_FC_surv(nn.Module):
def __init__(self, omic_input_dim=None, fusion=None, size_arg = "small", dropout=0.25, n_classes=4):
r"""
Attention MIL Implementation
Args:
omic_input_dim (int): Dimension size of genomic features.
fusion (str): Fusion method (Choices: concat, bilinear, or None)
size_arg (str): Size of NN architecture (Choices: small or large)
dropout (float): Dropout rate
n_classes (int): Output shape of NN
"""
super(MIL_Attention_FC_surv, self).__init__()
self.fusion = fusion
self.size_dict_path = {"small": [1024, 512, 256], "big": [1024, 512, 384]}
self.size_dict_omic = {'small': [256, 256]}
### Deep Sets Architecture Construction
size = self.size_dict_path[size_arg]
fc = [nn.Linear(size[0], size[1]), nn.ReLU(), nn.Dropout(dropout)]
attention_net = Attn_Net_Gated(L=size[1], D=size[2], dropout=dropout, n_classes=1)
fc.append(attention_net)
self.attention_net = nn.Sequential(*fc)
self.rho = nn.Sequential(*[nn.Linear(size[1], size[2]), nn.ReLU(), nn.Dropout(dropout)])
### Constructing Genomic SNN
if self.fusion is not None:
hidden = [256, 256]
fc_omic = [SNN_Block(dim1=omic_input_dim, dim2=hidden[0])]
for i, _ in enumerate(hidden[1:]):
fc_omic.append(SNN_Block(dim1=hidden[i], dim2=hidden[i+1], dropout=0.25))
self.fc_omic = nn.Sequential(*fc_omic)
if self.fusion == 'concat':
self.mm = nn.Sequential(*[nn.Linear(256*2, size[2]), nn.ReLU(), nn.Linear(size[2], size[2]), nn.ReLU()])
elif self.fusion == 'bilinear':
self.mm = BilinearFusion(dim1=256, dim2=256, scale_dim1=8, scale_dim2=8, mmhid=256)
else:
self.mm = None
self.classifier = nn.Linear(size[2], n_classes)
def relocate(self):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() >= 1:
device_ids = list(range(torch.cuda.device_count()))
self.attention_net = nn.DataParallel(self.attention_net, device_ids=device_ids).to('cuda:0')
if self.fusion is not None:
self.fc_omic = self.fc_omic.to(device)
self.mm = self.mm.to(device)
self.rho = self.rho.to(device)
self.classifier = self.classifier.to(device)
def forward(self, **kwargs):
x_path = kwargs['x_path']
A, h_path = self.attention_net(x_path)
A = torch.transpose(A, 1, 0)
A_raw = A
A = F.softmax(A, dim=1)
h_path = torch.mm(A, h_path)
h_path = self.rho(h_path).squeeze()
if self.fusion is not None:
x_omic = kwargs['x_omic']
h_omic = self.fc_omic(x_omic)
if self.fusion == 'bilinear':
h = self.mm(h_path.unsqueeze(dim=0), h_omic.unsqueeze(dim=0)).squeeze()
elif self.fusion == 'concat':
h = self.mm(torch.cat([h_path, h_omic], axis=0))
else:
h = h_path # [256] vector
logits = self.classifier(h).unsqueeze(0) # logits needs to be a [1 x 4] vector
Y_hat = torch.topk(logits, 1, dim = 1)[1]
hazards = torch.sigmoid(logits)
S = torch.cumprod(1 - hazards, dim=1)
return hazards, S, Y_hat, None, None
######################################
# Deep Attention MISL Implementation #
######################################
class MIL_Cluster_FC_surv(nn.Module):
def __init__(self, omic_input_dim=None, fusion=None, num_clusters=10, size_arg = "small", dropout=0.25, n_classes=4):
r"""
Attention MIL Implementation
Args:
omic_input_dim (int): Dimension size of genomic features.
fusion (str): Fusion method (Choices: concat, bilinear, or None)
size_arg (str): Size of NN architecture (Choices: small or large)
dropout (float): Dropout rate
n_classes (int): Output shape of NN
"""
super(MIL_Cluster_FC_surv, self).__init__()
self.size_dict_path = {"small": [1024, 512, 256], "big": [1024, 512, 384]}
self.size_dict_omic = {'small': [256, 256]}
self.num_clusters = num_clusters
self.fusion = fusion
### FC Cluster layers + Pooling
size = self.size_dict_path[size_arg]
phis = []
for phenotype_i in range(num_clusters):
phi = [nn.Linear(size[0], size[1]), nn.ReLU(), nn.Dropout(dropout),
nn.Linear(size[1], size[1]), nn.ReLU(), nn.Dropout(dropout)]
phis.append(nn.Sequential(*phi))
self.phis = nn.ModuleList(phis)
self.pool1d = nn.AdaptiveAvgPool1d(1)
### WSI Attention MIL Construction
fc = [nn.Linear(size[1], size[1]), nn.ReLU(), nn.Dropout(dropout)]
attention_net = Attn_Net_Gated(L=size[1], D=size[2], dropout=dropout, n_classes=1)
fc.append(attention_net)
self.attention_net = nn.Sequential(*fc)
self.rho = nn.Sequential(*[nn.Linear(size[1], size[2]), nn.ReLU(), nn.Dropout(dropout)])
### Genomic SNN Construction + Multimodal Fusion
if fusion is not None:
hidden = self.size_dict_omic['small']
fc_omic = [SNN_Block(dim1=omic_input_dim, dim2=hidden[0])]
for i, _ in enumerate(hidden[1:]):
fc_omic.append(SNN_Block(dim1=hidden[i], dim2=hidden[i+1], dropout=0.25))
self.fc_omic = nn.Sequential(*fc_omic)
if fusion == 'concat':
self.mm = nn.Sequential(*[nn.Linear(size[2]*2, size[2]), nn.ReLU(), nn.Linear(size[2], size[2]), nn.ReLU()])
elif self.fusion == 'bilinear':
self.mm = BilinearFusion(dim1=256, dim2=256, scale_dim1=8, scale_dim2=8, mmhid=256)
else:
self.mm = None
self.classifier = nn.Linear(size[2], n_classes)
def relocate(self):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() >= 1:
device_ids = list(range(torch.cuda.device_count()))
self.attention_net = nn.DataParallel(self.attention_net, device_ids=device_ids).to('cuda:0')
else:
self.attention_net = self.attention_net.to(device)
if self.fusion is not None:
self.fc_omic = self.fc_omic.to(device)
self.mm = self.mm.to(device)
self.phis = self.phis.to(device)
self.pool1d = self.pool1d.to(device)
self.rho = self.rho.to(device)
self.classifier = self.classifier.to(device)
def forward(self, **kwargs):
x_path = kwargs['x_path']
cluster_id = kwargs['cluster_id'].detach().cpu().numpy()
### FC Cluster layers + Pooling
h_cluster = []
for i in range(self.num_clusters):
h_cluster_i = self.phis[i](x_path[cluster_id==i])
if h_cluster_i.shape[0] == 0:
h_cluster_i = torch.zeros((1,512)).to(torch.device('cuda'))
h_cluster.append(self.pool1d(h_cluster_i.T.unsqueeze(0)).squeeze(2))
h_cluster = torch.stack(h_cluster, dim=1).squeeze(0)
### Attention MIL
A, h_path = self.attention_net(h_cluster)
A = torch.transpose(A, 1, 0)
A_raw = A
A = F.softmax(A, dim=1)
h_path = torch.mm(A, h_path)
h_path = self.rho(h_path).squeeze()
### Attention MIL + Genomic Fusion
if self.fusion is not None:
x_omic = kwargs['x_omic']
h_omic = self.fc_omic(x_omic)
if self.fusion == 'bilinear':
h = self.mm(h_path.unsqueeze(dim=0), h_omic.unsqueeze(dim=0)).squeeze()
elif self.fusion == 'concat':
h = self.mm(torch.cat([h_path, h_omic], axis=0))
else:
h = h_path
logits = self.classifier(h).unsqueeze(0)
Y_hat = torch.topk(logits, 1, dim = 1)[1]
hazards = torch.sigmoid(logits)
S = torch.cumprod(1 - hazards, dim=1)
return hazards, S, Y_hat, None, None
Datasets
Models
