LinearModule
This commit is contained in:
@@ -3,9 +3,9 @@ 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.blocks import ConvModule, LinearModule
|
||||
from ml_lib.modules.utils import (LightningBaseModule, HorizontalSplitter,
|
||||
HorizontalMerger)
|
||||
HorizontalMerger, F_x)
|
||||
from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
|
||||
BaseDataloadersMixin)
|
||||
|
||||
@@ -54,15 +54,11 @@ class BandwiseConvClassifier(BinaryMaskDatasetFunction,
|
||||
|
||||
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_1 = LinearModule(self.flat.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_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()
|
||||
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)
|
||||
@@ -74,13 +70,8 @@ class BandwiseConvClassifier(BinaryMaskDatasetFunction,
|
||||
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_3(tensor)
|
||||
tensor = self.full_out(tensor)
|
||||
tensor = self.sigmoid(tensor)
|
||||
return Namespace(main_out=tensor)
|
||||
|
||||
Reference in New Issue
Block a user