MILPLE/model.py at main · marrlab/MILPLE · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import os
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import numpy as np


class SCEMILA(nn.Module):

    def __init__(self, class_count, multi_attention, device):
        '''Initialize model.

        Accepts:
        - class_count (integer):
            amount of classes --> relevant for output vector
        - multi_attention (boolean):
            Defines if multiple attention values should be used for each image, one for each
            possible bag label.
        - device (String):
            either 'cuda:0' or the corresponding cpu counterpart.

        Returns:
        - Nothing, just set up model


        '''

        super(SCEMILA, self).__init__()

        # condense every image into self.L features (further encoding before actual MIL starts)
        self.L = 500
        self.D = 128                    # hidden layer size for attention network

        self.class_count = class_count
        self.multi_attention = multi_attention
        self.device = device


        # feature extractor before multiple instance learning starts
        self.FT_DIM_IN = 512
        self.ftr_proc = nn.Sequential(
            nn.Conv2d(self.FT_DIM_IN, int(self.FT_DIM_IN*1.5), kernel_size=2),
            nn.ReLU(),
            nn.Conv2d(int(self.FT_DIM_IN*1.5),
                      int(self.FT_DIM_IN*2), kernel_size=2),
            nn.ReLU(),
            nn.AdaptiveMaxPool2d(output_size=(1, 1)),
            nn.Flatten(),
            nn.Linear(int(self.FT_DIM_IN*2), self.L),
            nn.ReLU(),
        )

        # Networks for single attention approach
        # attention network (single attention approach)
        self.attention = nn.Sequential(
            nn.Linear(self.L, self.D),
            nn.Tanh(),
            nn.Linear(self.D, 1)
        )


        # classifier (single attention approach)
        self.classifier = nn.Sequential(
            nn.Linear(self.L, 64),
            nn.ReLU(),
            nn.Linear(64, self.class_count)
        )


    def forward(self, x):
        '''Forward pass of feature bag x through network.

        Accepts:
        - x (torch.Tensor):
            features of every single cell image from that patient

        Returns:
        - prediction (torch.Tensor):
            Activations of output layer
        - att_raw (torch.Tensor)
            Attention values before softmax transform
        - att_softmax (torch.Tensor)
            Attention values after softmax transform
        - bag_features (torch.Tensor)
            Bag features after attention calculation and matrix multiplication
        '''

        ft = x
        ft = self.ftr_proc(ft)

        # switch between multi- and single attention classification
        if(self.multi_attention):
            prediction = []
            bag_feature_stack = []
            attention_stack = []
            # calculate attention
            att_raw = self.attention_multi_column(ft)
            att_raw = torch.transpose(att_raw, 1, 0)

            # for every possible class, repeat
            for a in range(self.class_count):
                # softmax + Matrix multiplication
                att_softmax = F.softmax(att_raw[a, ...][None, ...], dim=1)
                bag_features = torch.mm(att_softmax, ft)
                bag_feature_stack.append(bag_features)
                # final classification with one output value (value indicating this specific class to be predicted)
                pred = self.classifier_multi_column[a](bag_features)
                prediction.append(pred)

            prediction = torch.stack(prediction).view(1, self.class_count)
            bag_feature_stack = torch.stack(bag_feature_stack).squeeze()
            # final softmax to obtain probabilities over all classes

            return prediction, att_raw, F.softmax(att_raw, dim=1), bag_feature_stack
        else:
            # calculate attention
            att_raw = self.attention(ft)
            att_raw = torch.transpose(att_raw, 1, 0)
            # Softmax + Matrix multiplication
            att_softmax = F.softmax(att_raw, dim=1)
            bag_features = torch.mm(att_softmax, ft)
            # final classification
            prediction = self.classifier(bag_features)


            return prediction