def test_layers_frozen(self):
with self.subTest('should_be_frozen'):
self._check_frozen(self.net.encoder)
with self.subTest('should_be_unfrozen'):
self.net.freeze_encoder = False
self.net.encoder = encoders.DenseEncoder((1, 32, 32), 3, 64)
self.net.train()
self._check_frozen(self.net.encoder, should_be_frozen=False)
def _build_networks(model_type, input_shape, latent_dim):
enc_dim = dec_dim = latent_dim
if model_type == 'vae' or model_type.startswith('beta_vae'):
enc_dim *= 2
num_layers = 3
if model_type == 'shallow':
encoder = encoders.ShallowEncoder(input_shape, enc_dim)
decoder = decoders.ShallowDecoder(dec_dim, input_shape)
elif model_type == 'stacked':
encoder = encoders.StackedEncoder(input_shape, num_layers, enc_dim)
decoder = decoders.StackedDecoder(dec_dim, num_layers, input_shape)
else:
encoder = encoders.DenseEncoder(input_shape, num_layers, enc_dim)
decoder = decoders.DenseDecoder(dec_dim, num_layers, input_shape)
return encoder, decoder
def setUp(self):
self.test_inputs = torch.randn(16, 1, 32, 32)
self.output_shape = torch.Size((16, 32))
self.net = encoders.DenseEncoder(self.test_inputs.shape[1:], 3, 32)
def setUp(self):
encoder = encoders.DenseEncoder((1, 32, 32), 3, 64)
bottleneck = bottlenecks.VariationalBottleneck(32)
self.net = Classifier(encoder, bottleneck, 10)
self.test_inputs = torch.randn(16, 1, 32, 32)
self.output_shape = torch.Size((16, 10))
def setUp(self):
self.enc = encoders.DenseEncoder((512, ), 2, 32)
self.dec = decoders.DenseDecoder(32, 2, (512, ))
self.neck = bottlenecks.IdentityBottleneck(32)
self.test_batch = [torch.randn(10, 512), torch.ones(10)]