requirements

This commit is contained in:
Si11ium
2020-05-14 23:08:36 +02:00
parent 407df15bbf
commit e7d1a4895a
9 changed files with 52 additions and 38 deletions

View File

@@ -5,11 +5,11 @@ from torch.nn import ModuleList
from ml_lib.modules.blocks import ConvModule, LinearModule, ResidualModule
from ml_lib.modules.utils import LightningBaseModule
from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetMixin,
BaseDataloadersMixin)
class ResidualConvClassifier(BinaryMaskDatasetFunction,
class ResidualConvClassifier(BinaryMaskDatasetMixin,
BaseDataloadersMixin,
BaseTrainMixin,
BaseValMixin,
@@ -45,6 +45,8 @@ class ResidualConvClassifier(BinaryMaskDatasetFunction,
last_shape = self.conv_list[-1].shape
self.conv_list.append(ConvModule(last_shape, filters, (k, k), conv_stride=(2, 2), conv_padding=2,
**self.params.module_kwargs))
for param in self.conv_list[-1].parameters():
param.requires_grad = False
last_shape = self.conv_list[-1].shape
self.full_1 = LinearModule(self.conv_list[-1].shape, self.params.lat_dim, **self.params.module_kwargs)