Parameter Adjustmens and Ensemble Model Implementation

models/bandwise_conv_classifier.py

@@ -0,0 +1,86 @@
+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.utils import (LightningBaseModule, HorizontalSplitter,
+                                  HorizontalMerger)
+from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
+                                BaseDataloadersMixin)
+
+
+class BandwiseConvClassifier(BinaryMaskDatasetFunction,
+                             BaseDataloadersMixin,
+                             BaseTrainMixin,
+                             BaseValMixin,
+                             BaseOptimizerMixin,
+                             LightningBaseModule
+                             ):
+
+    def __init__(self, hparams):
+        super(BandwiseConvClassifier, self).__init__(hparams)
+
+        # Dataset
+        # =============================================================================
+        self.dataset = self.build_dataset()
+
+        # Model Paramters
+        # =============================================================================
+        # Additional parameters
+        self.in_shape = self.dataset.train_dataset.sample_shape
+        self.conv_filters = self.params.filters
+        self.criterion = nn.BCELoss()
+        self.n_band_sections = 4
+
+        # 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.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_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()
+
+    def forward(self, batch, **kwargs):
+        tensors = self.split(batch)
+        for idx, tensor in enumerate(tensors):
+            tensors[idx] = self.conv_dict[f"conv_1_{idx}"](tensor)
+        for idx, tensor in enumerate(tensors):
+            tensors[idx] = self.conv_dict[f"conv_2_{idx}"](tensor)
+        for idx, tensor in enumerate(tensors):
+            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_out(tensor)
+        tensor = self.sigmoid(tensor)
+        return Namespace(main_out=tensor)