ResidualModule and New Parameters, Speed Manipulation
This commit is contained in:
@@ -1,11 +1,10 @@
|
||||
from argparse import Namespace
|
||||
|
||||
from torch import nn
|
||||
from torch.nn import ModuleDict, ModuleList
|
||||
from torch.nn import ModuleList
|
||||
|
||||
from ml_lib.modules.blocks import ConvModule, LinearModule
|
||||
from ml_lib.modules.utils import (LightningBaseModule, HorizontalSplitter,
|
||||
HorizontalMerger, F_x)
|
||||
from ml_lib.modules.utils import (LightningBaseModule, HorizontalSplitter, HorizontalMerger)
|
||||
from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
|
||||
BaseDataloadersMixin)
|
||||
|
||||
@@ -30,44 +29,39 @@ class BandwiseConvClassifier(BinaryMaskDatasetFunction,
|
||||
# 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)}
|
||||
)
|
||||
k = 3
|
||||
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), 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 = HorizontalMerger(self.conv_dict['conv_3_1'].shape, self.n_band_sections)
|
||||
self.merge = HorizontalMerger(self.band_list[-1][-1].shape, self.n_band_sections)
|
||||
|
||||
self.full_1 = LinearModule(self.flat.shape, self.params.lat_dim, **self.params.module_kwargs)
|
||||
self.full_1 = LinearModule(self.merge.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_3 = LinearModule(self.full_2.shape, self.full_2.shape // 2, **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):
|
||||
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)
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user