New Model, Many Changes

This commit is contained in:
Si11ium
2020-11-21 09:28:26 +01:00
parent 7bac9e984b
commit be097a111a
12 changed files with 349 additions and 125 deletions

View File

@@ -5,7 +5,7 @@ from torch import nn
from torch.nn import ModuleList
from ml_lib.modules.blocks import ConvModule, LinearModule
from ml_lib.modules.util import (LightningBaseModule, Flatten, HorizontalSplitter)
from ml_lib.modules.util import (LightningBaseModule, Splitter)
from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetMixin,
BaseDataloadersMixin)
@@ -69,7 +69,7 @@ class BandwiseConvMultiheadClassifier(BinaryMaskDatasetMixin,
# Modules
# =============================================================================
self.split = HorizontalSplitter(self.in_shape, self.n_band_sections)
self.split = Splitter(self.in_shape, self.n_band_sections)
self.band_list = ModuleList()
for band in range(self.n_band_sections):