Parameter Adjustmens and Ensemble Model Implementation

models/bandwise_conv_multihead_classifier.py

@@ -0,0 +1,107 @@
+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 ml_lib.modules.blocks import ConvModule
+from ml_lib.modules.utils import (LightningBaseModule, Flatten, HorizontalSplitter)
+from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
+                                BaseDataloadersMixin)
+
+
+class BandwiseConvMultiheadClassifier(BinaryMaskDatasetFunction,
+                                      BaseDataloadersMixin,
+                                      BaseTrainMixin,
+                                      BaseValMixin,
+                                      BaseOptimizerMixin,
+                                      LightningBaseModule
+                                      ):
+
+    def training_step(self, batch_xy, batch_nb, *args, **kwargs):
+        batch_x, batch_y = batch_xy
+        y = self(batch_x)
+        y, bands_y = y.main_out, y.bands
+        bands_y_losses = [self.criterion(band_y, batch_y) for band_y in bands_y]
+        return_dict = {f'band_{band_idx}_loss': band_y for band_idx, band_y in enumerate(bands_y_losses)}
+        overall_loss = self.criterion(y, batch_y)
+        combined_loss = overall_loss + torch.stack(bands_y_losses).sum()
+        return_dict.update(loss=combined_loss, overall_loss=overall_loss)
+        return return_dict
+
+    def validation_step(self, batch_xy, batch_idx, *args, **kwargs):
+        batch_x, batch_y = batch_xy
+        y = self(batch_x)
+        y, bands_y = y.main_out, y.bands
+        bands_y_losses = [self.criterion(band_y, batch_y) for band_y in bands_y]
+        return_dict = {f'band_{band_idx}_val_loss': band_y for band_idx, band_y in enumerate(bands_y_losses)}
+        overall_loss = self.criterion(y, batch_y)
+        combined_loss = overall_loss + torch.stack(bands_y_losses).sum()
+
+        val_abs_loss = self.absolute_loss(y, batch_y)
+        return_dict.update(val_bce_loss=combined_loss, val_abs_loss=val_abs_loss,
+                           batch_idx=batch_idx, y=y, batch_y=batch_y
+                           )
+        return return_dict
+
+    def __init__(self, hparams):
+        super(BandwiseConvMultiheadClassifier, 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 = 8
+
+        # 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.flat = Flatten(self.conv_dict['conv_3_1'].shape)
+        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)])
+
+        self.full_out = nn.Linear(self.n_band_sections, 1, self.params.bias)
+
+        # Utility Modules
+        self.sigmoid = 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)
+            tensor = self.bandwise_latent_list[idx](tensor)
+            tensor = self.bandwise_classifier_list[idx](tensor)
+            tensors[idx] = self.sigmoid(tensor)
+        tensor = torch.cat(tensors, dim=1)
+        tensor = self.full_out(tensor)
+        tensor = self.sigmoid(tensor)
+        return Namespace(main_out=tensor, bands=tensors)