requirements
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user