|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
def exists(val): |
|
return val is not None |
|
|
|
def initialize_weights(module): |
|
for m in module.modules(): |
|
if isinstance(m, nn.Linear): |
|
nn.init.xavier_normal_(m.weight) |
|
if m.bias is not None: |
|
m.bias.data.zero_() |
|
|
|
""" |
|
Attention Network with Sigmoid Gating (3 fc layers) |
|
args: |
|
L: input feature dimension |
|
D: hidden layer dimension |
|
dropout: whether to use dropout (p = 0.25) |
|
n_classes: number of classes |
|
""" |
|
class Attn_Net_Gated(nn.Module): |
|
|
|
def __init__(self, L = 1024, D = 256, n_tasks = 1): |
|
super(Attn_Net_Gated, self).__init__() |
|
self.attention_a = nn.Sequential(nn.Linear(L, D), nn.Tanh(), nn.Dropout(0.25)) |
|
self.attention_b = nn.Sequential(nn.Linear(L, D), nn.Sigmoid(), nn.Dropout(0.25)) |
|
self.attention_c = nn.Linear(D, n_tasks) |
|
|
|
def forward(self, x): |
|
a = self.attention_a(x) |
|
b = self.attention_b(x) |
|
A = a.mul(b) |
|
A = self.attention_c(A) |
|
return A, x |
|
|
|
|
|
""" |
|
Code borrow from: https://github.com/mahmoodlab/TOAD |
|
|
|
args: |
|
gate: whether to use gating in attention network |
|
size_args: size config of attention network |
|
dropout: whether to use dropout in attention network |
|
n_classes: number of classes |
|
""" |
|
|
|
class DeepAttnMIL(nn.Module): |
|
|
|
def __init__(self, input_dim = 1024, size_arg = "big", n_classes = 2): |
|
super(DeepAttnMIL, self).__init__() |
|
self.size_dict = {"small": [input_dim, 512, 256], "big": [input_dim, 512, 384]} |
|
size = self.size_dict[size_arg] |
|
|
|
self.attention_net = nn.Sequential( |
|
nn.Linear(size[0], size[1]), |
|
nn.ReLU(), |
|
nn.Dropout(0.25), |
|
Attn_Net_Gated(L = size[1], D = size[2], n_tasks = 1)) |
|
|
|
self.classifier = nn.Linear(size[1], n_classes) |
|
|
|
initialize_weights(self) |
|
|
|
def forward(self, h, return_features=False, attention_only=False): |
|
A, h = self.attention_net(h) |
|
A = torch.transpose(A, 1, 0) |
|
if attention_only: |
|
return A[0] |
|
|
|
A = F.softmax(A, dim=1) |
|
M = torch.mm(A, h) |
|
|
|
if return_features: |
|
return M |
|
|
|
logits = self.classifier(M) |
|
|
|
return logits |
|
|
|
|
