adding more samples

Author: ritheshkumar95

modules.py

@@ -15,29 +15,39 @@ def to_scalar(arr):
 def weights_init(m):
     classname = m.__class__.__name__
     if classname.find('Conv') != -1:
-        nn.init.xavier_uniform_(m.weight.data)
-        m.bias.data.fill_(0)
+        try:
+            nn.init.xavier_uniform_(m.weight.data)
+            m.bias.data.fill_(0)
+        except AttributeError:
+            print("Skipping initialization of ", classname)
 
 
 class VAE(nn.Module):
     def __init__(self, input_dim, dim, z_dim):
         super().__init__()
         self.encoder = nn.Sequential(
             nn.Conv2d(input_dim, dim, 4, 2, 1),
+            nn.BatchNorm2d(dim),
             nn.ReLU(True),
             nn.Conv2d(dim, dim, 4, 2, 1),
+            nn.BatchNorm2d(dim),
             nn.ReLU(True),
             nn.Conv2d(dim, dim, 5, 1, 0),
+            nn.BatchNorm2d(dim),
             nn.ReLU(True),
             nn.Conv2d(dim, z_dim * 2, 3, 1, 0),
+            nn.BatchNorm2d(z_dim * 2)
         )
 
         self.decoder = nn.Sequential(
             nn.ConvTranspose2d(z_dim, dim, 3, 1, 0),
+            nn.BatchNorm2d(dim),
             nn.ReLU(True),
             nn.ConvTranspose2d(dim, dim, 5, 1, 0),
+            nn.BatchNorm2d(dim),
             nn.ReLU(True),
             nn.ConvTranspose2d(dim, dim, 4, 2, 1),
+            nn.BatchNorm2d(dim),
             nn.ReLU(True),
             nn.ConvTranspose2d(dim, input_dim, 4, 2, 1),
             nn.Tanh()
@@ -204,7 +214,7 @@ def forward(self, x_v, x_h, h):
 class GatedPixelCNN(nn.Module):
     def __init__(self, input_dim=256, dim=64, n_layers=15):
         super().__init__()
-        self.dim = 64
+        self.dim = dim
 
         # Create embedding layer to embed input
         self.embedding = nn.Embedding(input_dim, dim)
@@ -225,11 +235,13 @@ def __init__(self, input_dim=256, dim=64, n_layers=15):
 
         # Add the output layer
         self.output_conv = nn.Sequential(
-            nn.Conv2d(dim, dim, 1),
+            nn.Conv2d(dim, 512, 1),
             nn.ReLU(True),
-            nn.Conv2d(dim, input_dim, 1)
+            nn.Conv2d(512, input_dim, 1)
         )
 
+        self.apply(weights_init)
+
     def forward(self, x, label):
         shp = x.size() + (-1, )
         x = self.embedding(x.view(-1)).view(shp)  # (B, H, W, C)

pixelcnn_prior.py

@@ -11,16 +11,16 @@
 N_EPOCHS = 100
 PRINT_INTERVAL = 100
 ALWAYS_SAVE = True
-DATASET = 'MNIST'  # CIFAR10 | MNIST | FashionMNIST
+DATASET = 'CIFAR10'  # CIFAR10 | MNIST | FashionMNIST
 NUM_WORKERS = 4
 
-LATENT_SHAPE = (7, 7)  # (8, 8) -> 32x32 images, (7, 7) -> 28x28 images
-INPUT_DIM = 1  # 3 (RGB) | 1 (Grayscale)
-DIM = 64
+LATENT_SHAPE = (8, 8)  # (8, 8) -> 32x32 images, (7, 7) -> 28x28 images
+INPUT_DIM = 3  # 3 (RGB) | 1 (Grayscale)
+DIM = 256
 VAE_DIM = 256
-N_LAYERS = 15
+N_LAYERS = 12
 K = 512
-LR = 1e-3
+LR = 3e-4
 
 DEVICE = torch.device('cuda') # torch.device('cpu')