dataset stretch now optional

This commit is contained in:
Si11ium
2020-05-15 09:43:30 +02:00
parent f458660613
commit 09f9dd9131
4 changed files with 13 additions and 9 deletions

View File

@@ -39,15 +39,18 @@ class ResidualConvClassifier(BinaryMaskDatasetMixin,
self.conv_list.append(ConvModule(last_shape, self.conv_filters[0], (k, k), conv_stride=(2, 2), conv_padding=1,
**self.params.module_kwargs))
last_shape = self.conv_list[-1].shape
for filters in self.conv_filters:
conv_module_params.update(conv_filters=filters)
for idx in range(len(self.conv_filters)):
conv_module_params.update(conv_filters=self.conv_filters[idx])
self.conv_list.append(ResidualModule(last_shape, ConvModule, 3, **conv_module_params))
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
try:
self.conv_list.append(ConvModule(last_shape, self.conv_filters[idx+1], (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
except IndexError:
pass
self.full_1 = LinearModule(self.conv_list[-1].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)