| 5 | import math |
| 6 | |
| 7 | class AttnNet(nn.Module): |
| 8 | # Adapted from https://github.com/mahmoodlab/CLAM/blob/master/models/model_clam.py |
| 9 | # Lu, M.Y., Williamson, D.F.K., Chen, T.Y. et al. Data-efficient and weakly supervised computational pathology on whole-slide images. Nat Biomed Eng 5, 555–570 (2021). https://doi.org/10.1038/s41551-020-00682-w |
| 10 | |
| 11 | def __init__(self, L=1024, D=256, dropout=False, p_dropout_atn=0.25, n_classes=1): |
| 12 | super(AttnNet, self).__init__() |
| 13 | |
| 14 | self.attention_a = [nn.Linear(L, D), nn.Tanh()] |
| 15 | |
| 16 | self.attention_b = [nn.Linear(L, D), nn.Sigmoid()] |
| 17 | |
| 18 | if dropout: |
| 19 | self.attention_a.append(nn.Dropout(p_dropout_atn)) |
| 20 | self.attention_b.append(nn.Dropout(p_dropout_atn)) |
| 21 | |
| 22 | self.attention_a = nn.Sequential(*self.attention_a) |
| 23 | self.attention_b = nn.Sequential(*self.attention_b) |
| 24 | self.attention_c = nn.Linear(D, n_classes) |
| 25 | |
| 26 | def forward(self, x): |
| 27 | a = self.attention_a(x) |
| 28 | b = self.attention_b(x) |
| 29 | A = a.mul(b) |
| 30 | A = self.attention_c(A) # N x n_classes |
| 31 | return A |
| 32 | |
| 33 | class Attn_Modality_Gated(nn.Module): |
| 34 | # Adapted from https://github.com/mahmoodlab/PORPOISE |