Hello @ptrblck,
First of all, thanks for your reply! I checked if any gradients were none, infinity or nan. Turns out that’s not the case still the model doesn’t train. But when I create a module with encoder and decoder in the module as submodules, it magically (for me) starts training as expected and produces really nice samples!! But I don’t get it. Shouldn’t the model train even if I have encoder and decoder as seperate modules and not shoved in nn.Module class? I may be wrong there as I have just started using torch.
My modified code:
Modified model
class VAE(nn.Module):
def __init__(self):
super(VAE, self).__init__()
self.encoder = encoder = nn.Sequential(OrderedDict([
('e_conv_layer_1', nn.Conv2d(1, 16, 5, 1)), # 16 x 24 x 24
('e_relu_layer_1', nn.LeakyReLU(inplace=True)),
('e_batch_norm_1', nn.BatchNorm2d(16)),
('e_conv_layer_2', nn.Conv2d(16, 32, 5, 1)), # 32 x 20 x 20
('e_relu_layer_2', nn.LeakyReLU(inplace=True)),
('e_batch_norm_2', nn.BatchNorm2d(32)),
('e_conv_layer_3', nn.Conv2d(32, 32, 11, 1)), # 32 x 10 x 10
('e_relu_layer_3', nn.LeakyReLU(inplace=True)),
('e_batch_norm_3', nn.BatchNorm2d(32)),
('e_conv_layer_4', nn.Conv2d(32, 64, 5, 1)), # 64 x 6 x 6
('e_relu_layer_4', nn.LeakyReLU(inplace=True)),
('e_batch_norm_4', nn.BatchNorm2d(64)),
('e_dropout_layer_1', nn.Dropout2d(p=0.75)),
('e_conv_layer_5', nn.Conv2d(64, 128, 5, 1)), # 128 x 2 x 2
('e_relu_layer_5', nn.LeakyReLU(inplace=True)),
('e_batch_norm_5', nn.BatchNorm2d(128)),
('e_dropout_layer_1', nn.Dropout2d(p=0.85)),
('e_flatten_layer', nn.Flatten()),
('e_out_layer', LocLogvar(128*2*2, latent_dims))
]))
self.decoder = decoder = nn.Sequential(OrderedDict([
('inv_linear_layer_1', nn.Linear(latent_dims, 128*2*2)), # 128 * 2 * 2
('inv_relu_layer_5', nn.LeakyReLU(inplace=True)),
('inv_flatten_layer', Reshape(128, 2, 2)), # 128 x 2 x 2
('inv_conv_layer_5', nn.ConvTranspose2d(128, 64, 5, 1)), # 64 x 6 x 6
('inv_batch_norm_5', nn.BatchNorm2d(64)),
('inv_relu_layer_4', nn.LeakyReLU(inplace=True)),
('inv_conv_layer_4', nn.ConvTranspose2d(64, 32, 5, 1)), # 32 x 10 x 10
('inv_batch_norm_4', nn.BatchNorm2d(32)),
('inv_relu_layer_3', nn.LeakyReLU(inplace=True)),
('inv_conv_layer_3', nn.ConvTranspose2d(32, 32, 11, 1)), # 32 x 20 x 20
('inv_batch_norm_3', nn.BatchNorm2d(32)),
('inv_relu_layer_2', nn.LeakyReLU(inplace=True)),
('inv_conv_layer_2', nn.ConvTranspose2d(32, 16, 5, 1)), # 16 x 24 x 24
('inv_batch_norm_2', nn.BatchNorm2d(16)),
('inv_relu_layer_1', nn.LeakyReLU(inplace=True)),
('inv_conv_layer_1', nn.ConvTranspose2d(16, 1, 5, 1)), # 1 x 28 x 28
('inv_batch_norm_1', nn.BatchNorm2d(1)),
('inv_out_layer', nn.Sigmoid())
]))
def encode(self, x):
return self.encoder(x)
def reparameterize(self, mu, logvar):
std = torch.exp(0.5*logvar)
eps = torch.randn_like(std)
return mu + eps*std
def decode(self, z):
return self.decoder(z)
def forward(self, x):
mu, logvar = self.encode(x)
z = self.reparameterize(mu, logvar)
return self.decode(z), mu, logvar