paper preperations and notebooks, optuna callbacks

models/bandwise_conv_classifier.py

@@ -0,0 +1,69 @@
+import inspect
+from argparse import Namespace
+
+from torch import nn
+from torch.nn import ModuleList
+
+from ml_lib.modules.blocks import ConvModule, LinearModule
+from ml_lib.modules.util import (LightningBaseModule, Splitter, Merger)
+from util.module_mixins import CombinedModelMixins
+
+
+class BandwiseConvClassifier(CombinedModelMixins,
+                             LightningBaseModule
+                             ):
+    def __init__(self, in_shape, n_classes, weight_init, activation,
+                 use_bias, use_norm, dropout, lat_dim, filters,
+                 lr, weight_decay, sto_weight_avg, lr_warm_restart_epochs, opt_reset_interval,
+                 loss, scheduler, lr_scheduler_parameter
+                 ):
+        # TODO: Move this to parent class, or make it much easieer to access....
+        a = dict(locals())
+        params = {arg: a[arg] for arg in inspect.signature(self.__init__).parameters.keys() if arg != 'self'}
+        super(BandwiseConvClassifier, self).__init__(params)
+
+        # Model Paramters
+        # =============================================================================
+        # Additional parameters
+        self.n_band_sections = 8
+
+        # Modules
+        # =============================================================================
+        self.split = Splitter(in_shape, self.n_band_sections)
+
+        k = 3
+        self.band_list = ModuleList()
+        for band in range(self.n_band_sections):
+            last_shape = self.split.shape[band]
+            conv_list = ModuleList()
+            for conv_filters in self.params.filters:
+                conv_list.append(ConvModule(last_shape, conv_filters, (k, k), conv_stride=(2, 2), conv_padding=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.merge = Merger(self.band_list[-1][-1].shape, self.n_band_sections)
+
+        self.full_1 = LinearModule(self.merge.shape, self.params.lat_dim, **self.params.module_kwargs)
+        self.full_2 = LinearModule(self.full_1.shape, self.params.lat_dim, **self.params.module_kwargs)
+
+        # Make Decision between binary and Multiclass Classification
+        logits = n_classes if n_classes > 2 else 1
+        module_kwargs = self.params.module_kwargs
+        module_kwargs.update(activation=(nn.Softmax if logits > 1 else nn.Sigmoid))
+        self.full_out = LinearModule(self.full_2.shape, logits, **module_kwargs)
+
+    def forward(self, batch, **kwargs):
+        tensors = self.split(batch)
+        for idx, (tensor, convs) in enumerate(zip(tensors, self.band_list)):
+            for conv in convs:
+                tensor = conv(tensor)
+            tensors[idx] = tensor
+
+        tensor = self.merge(tensors)
+        tensor = self.full_1(tensor)
+        tensor = self.full_2(tensor)
+        tensor = self.full_out(tensor)
+        return Namespace(main_out=tensor)