76 lines
3.1 KiB
Python
76 lines
3.1 KiB
Python
from argparse import Namespace
|
|
|
|
from torch import nn
|
|
from torch.nn import ModuleList
|
|
|
|
from ml_lib.modules.blocks import ConvModule
|
|
from ml_lib.modules.utils import LightningBaseModule, Flatten
|
|
from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
|
|
BaseDataloadersMixin)
|
|
|
|
|
|
class ConvClassifier(BinaryMaskDatasetFunction,
|
|
BaseDataloadersMixin,
|
|
BaseTrainMixin,
|
|
BaseValMixin,
|
|
BaseOptimizerMixin,
|
|
LightningBaseModule
|
|
):
|
|
|
|
def __init__(self, hparams):
|
|
super(ConvClassifier, 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()
|
|
|
|
# Modules with Parameters
|
|
self.conv_list = ModuleList()
|
|
last_shape = self.in_shape
|
|
k = 3 # Base Kernel Value
|
|
for filters in self.conv_filters:
|
|
self.conv_list.append(ConvModule(last_shape, filters, (k, k*2), conv_stride=2, **self.params.module_kwargs))
|
|
last_shape = self.conv_list[-1].shape
|
|
self.conv_list.appen(ConvModule(last_shape, filters, 1, conv_stride=1, **self.params.module_kwargs))
|
|
last_shape = self.conv_list[-1].shape
|
|
self.conv_list.appen(ConvModule(last_shape, 1, 1, conv_stride=1, **self.params.module_kwargs))
|
|
last_shape = self.conv_list[-1].shape
|
|
k = k+2
|
|
|
|
self.flat = Flatten(self.conv_list[-1].shape)
|
|
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_3 = nn.Linear(self.full_2.out_features, self.full_2.out_features // 2, self.params.bias)
|
|
|
|
self.full_out = nn.Linear(self.full_3.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()
|
|
|
|
def forward(self, batch, **kwargs):
|
|
tensor = batch
|
|
for conv in self.conv_list:
|
|
tensor = conv(tensor)
|
|
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.activation(tensor)
|
|
tensor = self.dropout(tensor)
|
|
tensor = self.full_out(tensor)
|
|
tensor = self.sigmoid(tensor)
|
|
return Namespace(main_out=tensor)
|