LinearModule

This commit is contained in:
Si11ium
2020-05-09 21:56:58 +02:00
parent 5e6b0e598f
commit 3fbc98dfa3
7 changed files with 145 additions and 138 deletions

View File

@@ -3,8 +3,8 @@ 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 ml_lib.modules.blocks import ConvModule, LinearModule
from ml_lib.modules.utils import LightningBaseModule
from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
BaseDataloadersMixin)
@@ -38,38 +38,21 @@ class ConvClassifier(BinaryMaskDatasetFunction,
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.conv_list.append(ConvModule(last_shape, 1, 1, conv_stride=1, **self.params.module_kwargs))
# last_shape = self.conv_list[-1].shape
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_1 = LinearModule(self.flat.shape, self.params.lat_dim, **self.params.module_kwargs)
self.full_2 = LinearModule(self.full_1.out_features, self.full_1.out_features * 2, self.params.bias)
self.full_3 = LinearModule(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()
self.full_out = LinearModule(self.full_3.out_features, 1, bias=self.params.bias, activation=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)