| 33 | |
| 34 | class Im4MEC(nn.Module): |
| 35 | def __init__( |
| 36 | self, |
| 37 | input_feature_size=1024, |
| 38 | precompression_layer=True, |
| 39 | feature_size_comp = 512, |
| 40 | feature_size_attn = 256, |
| 41 | dropout=True, |
| 42 | p_dropout_fc=0.25, |
| 43 | p_dropout_atn=0.25, |
| 44 | n_classes=4, |
| 45 | ): |
| 46 | super(Im4MEC, self).__init__() |
| 47 | |
| 48 | self.n_classes = n_classes |
| 49 | |
| 50 | if precompression_layer: |
| 51 | self.compression_layer = nn.Sequential(*[ |
| 52 | nn.Linear(input_feature_size, feature_size_comp*4), |
| 53 | nn.ReLU(), |
| 54 | nn.Dropout(p_dropout_fc), |
| 55 | nn.Linear(feature_size_comp*4, feature_size_comp*2), |
| 56 | nn.ReLU(), |
| 57 | nn.Dropout(p_dropout_fc), |
| 58 | nn.Linear(feature_size_comp*2, feature_size_comp), |
| 59 | nn.ReLU(), |
| 60 | nn.Dropout(p_dropout_fc)]) |
| 61 | |
| 62 | dim_post_compression = feature_size_comp |
| 63 | else: |
| 64 | self.compression_layer = nn.Identity() |
| 65 | dim_post_compression = input_feature_size |
| 66 | |
| 67 | self.attention_net = Attn_Net_Gated( |
| 68 | L=dim_post_compression, |
| 69 | D=feature_size_attn, |
| 70 | dropout=dropout, |
| 71 | p_dropout_atn=p_dropout_atn, |
| 72 | n_classes=self.n_classes) |
| 73 | |
| 74 | # Classification head. |
| 75 | self.classifiers = nn.ModuleList( |
| 76 | [nn.Linear(dim_post_compression, 1) for i in range(self.n_classes)] |
| 77 | ) |
| 78 | |
| 79 | # Init weights. |
| 80 | self.apply(self._init_weights) |
| 81 | |
| 82 | def _init_weights(self, module): |
| 83 | if isinstance(module, nn.Linear): |