def test_decoder(self):
np.random.seed(42)
torch.manual_seed(42)
z_dim = 20
decoder = MLPDecoder(z_dim=20, output_shape=[1,28,28])
z = torch.randn(64, z_dim)
imgs = decoder(z)
self.assertTrue(len(imgs.shape) == 4 and all([imgs.shape[i] == o for i,o in enumerate([64,1,28,28])]),
msg="Output of the decoder should be an image with shape [B,C,H,W], but got: %s." % str(imgs.shape))
self.assertTrue((imgs < 0).any(),
msg="The output of the decoder does not have any negative values. " + \
"You might be applying a sigmoid on the decoder output. " + \
"It is recommended to work on logits instead as this is numercially more stable. " + \
"Ensure that you are using the correct loss accordings (BCEWithLogits instead of BCE).")
def __init__(self, model_name, hidden_dims, num_filters, z_dim, lr):
"""
PyTorch Lightning module that summarizes all components to train a VAE.
Inputs:
model_name - String denoting what encoder/decoder class to use. Either 'MLP' or 'CNN'
hidden_dims - List of hidden dimensionalities to use in the MLP layers of the encoder (decoder reversed)
num_filters - Number of channels to use in a CNN encoder/decoder
z_dim - Dimensionality of latent space
lr - Learning rate to use for the optimizer
"""
super().__init__()
self.save_hyperparameters()
self.bce_loss = nn.BCEWithLogitsLoss()
if model_name == 'MLP':
self.encoder = MLPEncoder(z_dim=z_dim, hidden_dims=hidden_dims)
self.decoder = MLPDecoder(z_dim=z_dim, hidden_dims=hidden_dims[::-1])
else:
self.encoder = CNNEncoder(z_dim=z_dim, num_filters=num_filters)
self.decoder = CNNDecoder(z_dim=z_dim, num_filters=num_filters)
def __init__(self, model_name, hidden_dims, num_filters, z_dim, *args,
**kwargs):
"""
PyTorch module that summarizes all components to train a VAE.
Inputs:
model_name - String denoting what encoder/decoder class to use. Either 'MLP' or 'CNN'
hidden_dims - List of hidden dimensionalities to use in the MLP layers of the encoder (decoder reversed)
num_filters - Number of channels to use in a CNN encoder/decoder
z_dim - Dimensionality of latent space
"""
super().__init__()
self.z_dim = z_dim
if model_name == 'MLP':
self.encoder = MLPEncoder(z_dim=z_dim, hidden_dims=hidden_dims)
self.decoder = MLPDecoder(z_dim=z_dim,
hidden_dims=hidden_dims[::-1])
else:
self.encoder = CNNEncoder(z_dim=z_dim, num_filters=num_filters)
self.decoder = CNNDecoder(z_dim=z_dim, num_filters=num_filters)