LinearModule

models/bandwise_conv_classifier.py

@@ -3,9 +3,9 @@ from argparse import Namespace
 from torch import nn
 from torch.nn import ModuleDict, ModuleList
 
-from ml_lib.modules.blocks import ConvModule
+from ml_lib.modules.blocks import ConvModule, LinearModule
 from ml_lib.modules.utils import (LightningBaseModule, HorizontalSplitter,
-                                  HorizontalMerger)
+                                  HorizontalMerger, F_x)
 from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
                                 BaseDataloadersMixin)
 
@@ -54,15 +54,11 @@ class BandwiseConvClassifier(BinaryMaskDatasetFunction,
 
         self.merge = HorizontalMerger(self.conv_dict['conv_3_1'].shape, self.n_band_sections)
 
-        self.full_1 = nn.Linear(self.flat.shape, self.params.lat_dim, self.params.bias)
-        self.full_2 = nn.Linear(self.full_1.out_features, self.full_1.out_features // 2, self.params.bias)
+        self.full_1 = LinearModule(self.flat.shape, self.params.lat_dim, **self.params.module_kwargs)
+        self.full_2 = LinearModule(self.full_1.shape, self.full_1.shape * 2, **self.params.module_kwargs)
+        self.full_3 = LinearModule(self.full_2.shape, self.full_2.out_features // 2, **self.params.module_kwargs)
 
-        self.full_out = nn.Linear(self.full_2.out_features, 1, self.params.bias)
-
-        # Utility Modules
-        self.dropout = nn.Dropout2d(self.params.dropout) if self.params.dropout else lambda x: x
-        self.activation = self.params.activation()
-        self.sigmoid = nn.Sigmoid()
+        self.full_out = LinearModule(self.full_3.shape, 1, bias=self.params.bias, activation=nn.Sigmoid)
 
     def forward(self, batch, **kwargs):
         tensors = self.split(batch)
@@ -74,13 +70,8 @@ class BandwiseConvClassifier(BinaryMaskDatasetFunction,
             tensors[idx] = self.conv_dict[f"conv_3_{idx}"](tensor)
 
         tensor = self.merge(tensors)
-        tensor = self.flat(tensor)
         tensor = self.full_1(tensor)
-        tensor = self.activation(tensor)
-        tensor = self.dropout(tensor)
         tensor = self.full_2(tensor)
-        tensor = self.activation(tensor)
-        tensor = self.dropout(tensor)
+        tensor = self.full_3(tensor)
         tensor = self.full_out(tensor)
-        tensor = self.sigmoid(tensor)
         return Namespace(main_out=tensor)

models/bandwise_conv_multihead_classifier.py

@@ -1,12 +1,10 @@
 from argparse import Namespace
-from collections import defaultdict
 
 import torch
 from torch import nn
-from torch.nn import ModuleDict, ModuleList
-from torchcontrib.optim import SWA
+from torch.nn import ModuleList
 
-from ml_lib.modules.blocks import ConvModule
+from ml_lib.modules.blocks import ConvModule, LinearModule
 from ml_lib.modules.utils import (LightningBaseModule, Flatten, HorizontalSplitter)
 from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
                                 BaseDataloadersMixin)
@@ -59,49 +57,57 @@ class BandwiseConvMultiheadClassifier(BinaryMaskDatasetFunction,
         self.in_shape = self.dataset.train_dataset.sample_shape
         self.conv_filters = self.params.filters
         self.criterion = nn.BCELoss()
-        self.n_band_sections = 8
+        self.n_band_sections = 4
+        k = 3  # Base Kernel Value
 
         # Modules
         # =============================================================================
         self.split = HorizontalSplitter(self.in_shape, self.n_band_sections)
-        self.conv_dict = ModuleDict()
 
-        self.conv_dict.update({f"conv_1_{band_section}":
-                                   ConvModule(self.split.shape, self.conv_filters[0], 3, conv_stride=1, **self.params.module_kwargs)
-                               for band_section in range(self.n_band_sections)}
-                              )
-        self.conv_dict.update({f"conv_2_{band_section}":
-                                   ConvModule(self.conv_dict['conv_1_1'].shape, self.conv_filters[1], 3, conv_stride=1,
-                                              **self.params.module_kwargs) for band_section in range(self.n_band_sections)}
-                              )
-        self.conv_dict.update({f"conv_3_{band_section}":
-                                   ConvModule(self.conv_dict['conv_2_1'].shape, self.conv_filters[2], 3, conv_stride=1,
-                                              **self.params.module_kwargs)
-                               for band_section in range(self.n_band_sections)}
-                              )
+        self.band_list = ModuleList()
+        for band in range(self.n_band_sections):
+            last_shape = self.split.shape
+            conv_list = ModuleList()
+            for filters in self.conv_filters:
+                conv_list.append(ConvModule(last_shape, filters, (k, k*4), conv_stride=(1, 2),
+                                            **self.params.module_kwargs))
+                last_shape = conv_list[-1].shape
+                # self.conv_list.append(ConvModule(last_shape, 1, 1, conv_stride=1, **self.params.module_kwargs))
+                # last_shape = self.conv_list[-1].shape
+            self.band_list.append(conv_list)
 
-        self.flat = Flatten(self.conv_dict['conv_3_1'].shape)
+        self.bandwise_deep_list_1 = ModuleList([
+            LinearModule(self.band_list[0][-1].shape, self.params.lat_dim * 4, **self.params.module_kwargs)
+            for _ in range(self.n_band_sections)])
+        self.bandwise_deep_list_2 = ModuleList([
+            LinearModule(self.params.lat_dim * 4, self.params.lat_dim * 2, **self.params.module_kwargs)
+            for _ in range(self.n_band_sections)])
         self.bandwise_latent_list = ModuleList([
-            nn.Linear(self.flat.shape, self.params.lat_dim, self.params.bias) for _ in range(self.n_band_sections)])
-        self.bandwise_classifier_list = ModuleList([nn.Linear(self.params.lat_dim, 1, self.params.bias)
-                                                    for _ in range(self.n_band_sections)])
+            LinearModule(self.params.lat_dim * 2, self.params.lat_dim, **self.params.module_kwargs)
+            for _ in range(self.n_band_sections)])
+        self.bandwise_classifier_list = ModuleList([
+            LinearModule(self.params.lat_dim, 1, bias=self.params.bias, activation=nn.Sigmoid)
+            for _ in range(self.n_band_sections)])
 
-        self.full_out = nn.Linear(self.n_band_sections, 1, self.params.bias)
-
-        # Utility Modules
-        self.sigmoid = nn.Sigmoid()
+        self.full_1 = LinearModule(self.n_band_sections, self.params.lat_dim * 4, **self.params.module_kwargs)
+        self.full_2 = LinearModule(self.full_1.shape, self.params.lat_dim * 2, **self.params.module_kwargs)
+        self.full_3 = LinearModule(self.full_2.shape, self.params.lat_dim, **self.params.module_kwargs)
+        self.full_out = LinearModule(self.full_3.shape, 1, bias=self.params.bias, activation=nn.Sigmoid)
 
     def forward(self, batch, **kwargs):
         tensors = self.split(batch)
-        for idx, tensor in enumerate(tensors):
-            tensor = self.conv_dict[f"conv_1_{idx}"](tensor)
-            tensor = self.conv_dict[f"conv_2_{idx}"](tensor)
-            tensor = self.conv_dict[f"conv_3_{idx}"](tensor)
-            tensor = self.flat(tensor)
+        for idx, (tensor, convs) in enumerate(zip(tensors, self.band_list)):
+            for conv in convs:
+                tensor = conv(tensor)
+
+            tensor = self.bandwise_deep_list_1[idx](tensor)
+            tensor = self.bandwise_deep_list_2[idx](tensor)
             tensor = self.bandwise_latent_list[idx](tensor)
-            tensor = self.bandwise_classifier_list[idx](tensor)
-            tensors[idx] = self.sigmoid(tensor)
+            tensors[idx] = self.bandwise_classifier_list[idx](tensor)
+
         tensor = torch.cat(tensors, dim=1)
+        tensor = self.full_1(tensor)
+        tensor = self.full_2(tensor)
+        tensor = self.full_3(tensor)
         tensor = self.full_out(tensor)
-        tensor = self.sigmoid(tensor)
         return Namespace(main_out=tensor, bands=tensors)

models/conv_classifier.py

@@ -3,8 +3,8 @@ from argparse import Namespace
 from torch import nn
 from torch.nn import ModuleList
 
-from ml_lib.modules.blocks import ConvModule
-from ml_lib.modules.utils import LightningBaseModule, Flatten
+from ml_lib.modules.blocks import ConvModule, LinearModule
+from ml_lib.modules.utils import LightningBaseModule
 from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
                                 BaseDataloadersMixin)
 
@@ -38,38 +38,21 @@ class ConvClassifier(BinaryMaskDatasetFunction,
         for filters in self.conv_filters:
             self.conv_list.append(ConvModule(last_shape, filters, (k, k*2), conv_stride=2, **self.params.module_kwargs))
             last_shape = self.conv_list[-1].shape
-            self.conv_list.appen(ConvModule(last_shape, filters, 1, conv_stride=1, **self.params.module_kwargs))
-            last_shape = self.conv_list[-1].shape
-            self.conv_list.appen(ConvModule(last_shape, 1, 1, conv_stride=1, **self.params.module_kwargs))
-            last_shape = self.conv_list[-1].shape
-            k = k+2
+            # self.conv_list.append(ConvModule(last_shape, 1, 1, conv_stride=1, **self.params.module_kwargs))
+            # last_shape = self.conv_list[-1].shape
 
-        self.flat = Flatten(self.conv_list[-1].shape)
-        self.full_1 = nn.Linear(self.flat.shape, self.params.lat_dim, self.params.bias)
-        self.full_2 = nn.Linear(self.full_1.out_features, self.full_1.out_features * 2, self.params.bias)
-        self.full_3 = nn.Linear(self.full_2.out_features, self.full_2.out_features // 2, self.params.bias)
+        self.full_1 = LinearModule(self.flat.shape, self.params.lat_dim, **self.params.module_kwargs)
+        self.full_2 = LinearModule(self.full_1.out_features, self.full_1.out_features * 2, self.params.bias)
+        self.full_3 = LinearModule(self.full_2.out_features, self.full_2.out_features // 2, self.params.bias)
 
-        self.full_out = nn.Linear(self.full_3.out_features, 1, self.params.bias)
-
-        # Utility Modules
-        self.dropout = nn.Dropout2d(self.params.dropout) if self.params.dropout else lambda x: x
-        self.activation = self.params.activation()
-        self.sigmoid = nn.Sigmoid()
+        self.full_out = LinearModule(self.full_3.out_features, 1, bias=self.params.bias, activation=nn.Sigmoid)
 
     def forward(self, batch, **kwargs):
         tensor = batch
         for conv in self.conv_list:
             tensor = conv(tensor)
-        tensor = self.flat(tensor)
         tensor = self.full_1(tensor)
-        tensor = self.activation(tensor)
-        tensor = self.dropout(tensor)
         tensor = self.full_2(tensor)
-        tensor = self.activation(tensor)
-        tensor = self.dropout(tensor)
         tensor = self.full_3(tensor)
-        tensor = self.activation(tensor)
-        tensor = self.dropout(tensor)
         tensor = self.full_out(tensor)
-        tensor = self.sigmoid(tensor)
         return Namespace(main_out=tensor)