making code neater

Author: ritheshkumar95

main.py

@@ -2,50 +2,52 @@
 import torch.nn.functional as F
 from torchvision import datasets, transforms
 from modules import AutoEncoder, to_scalar
-from torch.autograd import Variable
 import numpy as np
 from torchvision.utils import save_image
 import time
 
 
 BATCH_SIZE = 128
+N_EPOCHS = 100
+PRINT_INTERVAL = 100
+DATASET = 'FashionMNIST'  # CIFAR10 | MNIST | FashionMNIST
 NUM_WORKERS = 4
-LR = 2e-4
+
+INPUT_DIM = 1  # 3 (RGB) | 1 (Grayscale)
+DIM = 256
 K = 512
 LAMDA = 1
-PRINT_INTERVAL = 100
-N_EPOCHS = 100
+LR = 2e-4
 
 
 preproc_transform = transforms.Compose([
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 ])
 train_loader = torch.utils.data.DataLoader(
-    datasets.CIFAR10(
-        '../data/cifar10/', train=True, download=True,
+    eval('datasets.'+DATASET)(
+        '../data/{}/'.format(DATASET), train=True, download=True,
         transform=preproc_transform,
     ), batch_size=BATCH_SIZE, shuffle=False,
     num_workers=NUM_WORKERS, pin_memory=True
 )
-
 test_loader = torch.utils.data.DataLoader(
-    datasets.CIFAR10(
-        '../data/cifar10/', train=False,
+    eval('datasets.'+DATASET)(
+        '../data/{}/'.format(DATASET), train=False,
         transform=preproc_transform
     ), batch_size=BATCH_SIZE, shuffle=False,
     num_workers=NUM_WORKERS, pin_memory=True
 )
 
-model = AutoEncoder(K).cuda()
+model = AutoEncoder(INPUT_DIM, DIM, K).cuda()
 opt = torch.optim.Adam(model.parameters(), lr=LR)
 
 
 def train():
     train_loss = []
-    for batch_idx, (data, _) in enumerate(train_loader):
+    for batch_idx, (x, _) in enumerate(train_loader):
         start_time = time.time()
-        x = Variable(data, requires_grad=False).cuda()
+        x = x.cuda()
 
         opt.zero_grad()
 
@@ -70,20 +72,20 @@ def train():
 
         train_loss.append(to_scalar([loss_recons, loss_vq]))
 
-        if (batch_idx + 1) % 100 == 0:
+        if (batch_idx + 1) % PRINT_INTERVAL == 0:
             print('\tIter [{}/{} ({:.0f}%)]\tLoss: {} Time: {}'.format(
-                batch_idx * len(data), len(train_loader.dataset),
-                100. * batch_idx / len(train_loader),
-                np.asarray(train_loss)[-100:].mean(0),
+                batch_idx * len(x), len(train_loader.dataset),
+                PRINT_INTERVAL * batch_idx / len(train_loader),
+                np.asarray(train_loss)[-PRINT_INTERVAL:].mean(0),
                 time.time() - start_time
             ))
 
 
 def test():
     start_time = time.time()
     val_loss = []
-    for batch_idx, (data, _) in enumerate(test_loader):
-        x = Variable(data, volatile=True).cuda()
+    for batch_idx, (x, _) in enumerate(test_loader):
+        x = x.cuda()
         x_tilde, z_e_x, z_q_x = model(x)
         loss_recons = F.mse_loss(x_tilde, x)
         loss_vq = F.mse_loss(z_q_x, z_e_x.detach())
@@ -98,12 +100,17 @@ def test():
 
 def generate_samples():
     x, _ = test_loader.__iter__().next()
-    x = Variable(x[:32]).cuda()
+    x = x[:32].cuda()
     x_tilde, _, _ = model(x)
 
     x_cat = torch.cat([x, x_tilde], 0)
     images = (x_cat.cpu().data + 1) / 2
-    save_image(images, './sample_cifar.png', nrow=8)
+
+    save_image(
+        images,
+        'samples/reconstructions_{}.png'.format(DATASET),
+        nrow=8
+    )
 
 
 BEST_LOSS = 999
@@ -117,7 +124,7 @@ def generate_samples():
         BEST_LOSS = cur_loss
         LAST_SAVED = epoch
         print("Saving model!")
-        torch.save(model.state_dict(), 'best_autoencoder.pt')
+        torch.save(model.state_dict(), 'models/{}_autoencoder.pt'.format(DATASET))
     else:
         print("Not saving model! Last saved: {}".format(LAST_SAVED))
 

modules.py

@@ -1,7 +1,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.autograd import Variable
 
 
 def to_scalar(arr):
@@ -33,27 +32,27 @@ def forward(self, x):
 
 
 class AutoEncoder(nn.Module):
-    def __init__(self, K=512):
+    def __init__(self, input_dim, dim, K=512):
         super(AutoEncoder, self).__init__()
         self.encoder = nn.Sequential(
-            nn.Conv2d(3, 256, 4, 2, 1),
+            nn.Conv2d(input_dim, dim, 4, 2, 1),
             nn.ReLU(True),
-            nn.Conv2d(256, 256, 4, 2, 1),
-            ResBlock(256),
-            ResBlock(256),
+            nn.Conv2d(dim, dim, 4, 2, 1),
+            ResBlock(dim),
+            ResBlock(dim),
         )
 
-        self.embedding = nn.Embedding(K, 256)
+        self.embedding = nn.Embedding(K, dim)
         # self.embedding.weight.data.copy_(1./K * torch.randn(K, 256))
         self.embedding.weight.data.uniform_(-1./K, 1./K)
 
         self.decoder = nn.Sequential(
-            ResBlock(256),
-            ResBlock(256),
+            ResBlock(dim),
+            ResBlock(dim),
             nn.ReLU(True),
-            nn.ConvTranspose2d(256, 256, 4, 2, 1),
+            nn.ConvTranspose2d(dim, dim, 4, 2, 1),
             nn.ReLU(True),
-            nn.ConvTranspose2d(256, 3, 4, 2, 1),
+            nn.ConvTranspose2d(dim, input_dim, 4, 2, 1),
             nn.Tanh()
         )
 
@@ -94,12 +93,16 @@ def forward(self, x):
 
 
 class GatedMaskedConv2d(nn.Module):
-    def __init__(self, mask_type, dim, kernel, residual=True):
+    def __init__(self, mask_type, dim, kernel, residual=True, n_classes=10):
         super().__init__()
         assert kernel % 2 == 1, print("Kernel size must be odd")
         self.mask_type = mask_type
         self.residual = residual
 
+        self.class_cond_embedding = nn.Embedding(
+            n_classes, 2 * dim
+        )
+
         kernel_shp = (kernel // 2 + 1, kernel)  # (ceil(n/2), n)
         padding_shp = (kernel // 2, kernel // 2)
         self.vert_stack = nn.Conv2d(
@@ -124,19 +127,20 @@ def make_causal(self):
         self.vert_stack.weight.data[:, :, -1].zero_()  # Mask final row
         self.horiz_stack.weight.data[:, :, :, -1].zero_()  # Mask final column
 
-    def forward(self, x_v, x_h):
+    def forward(self, x_v, x_h, h):
         if self.mask_type == 'A':
             self.make_causal()
 
+        h = self.class_cond_embedding(h)
         h_vert = self.vert_stack(x_v)
         h_vert = h_vert[:, :, :x_v.size(-1), :]
-        out_v = self.gate(h_vert)
+        out_v = self.gate(h_vert + h[:, :, None, None])
 
         h_horiz = self.horiz_stack(x_h)
         h_horiz = h_horiz[:, :, :, :x_h.size(-2)]
         v2h = self.vert_to_horiz(h_vert)
 
-        out = self.gate(v2h + h_horiz)
+        out = self.gate(v2h + h_horiz + h[:, :, None, None])
         if self.residual:
             out_h = self.horiz_resid(out) + x_h
         else:
@@ -174,25 +178,23 @@ def __init__(self, input_dim=256, dim=64, n_layers=15):
             nn.Conv2d(dim, input_dim, 1)
         )
 
-    def forward(self, x):
+    def forward(self, x, label):
         shp = x.size() + (-1, )
         x = self.embedding(x.view(-1)).view(shp)  # (B, H, W, C)
         x = x.permute(0, 3, 1, 2)  # (B, C, W, W)
 
         x_v, x_h = (x, x)
         for i, layer in enumerate(self.layers):
-            x_v, x_h = layer(x_v, x_h)
+            x_v, x_h = layer(x_v, x_h, label)
 
         return self.output_conv(x_h)
 
-    def generate(self, batch_size=64):
-        x = Variable(
-            torch.zeros(64, 8, 8).long()
-        ).cuda()
+    def generate(self, label, shape=(8, 8), batch_size=64):
+        x = torch.zeros(batch_size, *shape).long().cuda()
 
-        for i in range(8):
-            for j in range(8):
-                logits = self.forward(x)
+        for i in range(shape[0]):
+            for j in range(shape[1]):
+                logits = self.forward(x, label)
                 probs = F.softmax(logits[:, :, i, j], -1)
                 x.data[:, i, j].copy_(
                     probs.multinomial(1).squeeze().data