improving code and adding pixelcnn

Author: ritheshkumar95

main.py

@@ -8,31 +8,40 @@
 import time
 
 
-kwargs = {'num_workers': 1, 'pin_memory': True}
+BATCH_SIZE = 128
+NUM_WORKERS = 4
+LR = 2e-4
+K = 256
+LAMDA = 0.25
+PRINT_INTERVAL = 100
+N_EPOCHS = 100
+
+
+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,
-        transform=transforms.Compose(
-            [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
-        )
-    ), batch_size=64, shuffle=False, **kwargs
+        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,
-        transform=transforms.Compose(
-            [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
-        )
-    ), batch_size=32, shuffle=False, **kwargs
+        transform=preproc_transform
+    ), batch_size=BATCH_SIZE, shuffle=False,
+    num_workers=NUM_WORKERS, pin_memory=True
 )
-test_data = list(test_loader)
 
-model = AutoEncoder().cuda()
-opt = torch.optim.Adam(model.parameters(), lr=3e-4)
+model = AutoEncoder(K).cuda()
+opt = torch.optim.Adam(model.parameters(), lr=LR)
 
 
-def train(epoch):
+def train():
     train_loss = []
     for batch_idx, (data, _) in enumerate(train_loader):
         start_time = time.time()
@@ -43,43 +52,75 @@ def train(epoch):
         x_tilde, z_e_x, z_q_x = model(x)
         z_q_x.retain_grad()
 
-        loss_recons = F.l1_loss(x_tilde, x)
+        loss_recons = F.mse_loss(x_tilde, x)
         loss_recons.backward(retain_graph=True)
 
         # Straight-through estimator
         z_e_x.backward(z_q_x.grad, retain_graph=True)
 
         # Vector quantization objective
+        model.embedding.zero_grad()
         loss_vq = F.mse_loss(z_q_x, z_e_x.detach())
         loss_vq.backward(retain_graph=True)
 
         # Commitment objective
-        loss_commit = 0.25 * F.mse_loss(z_e_x, z_q_x.detach())
+        loss_commit = LAMDA * F.mse_loss(z_e_x, z_q_x.detach())
         loss_commit.backward()
         opt.step()
 
         train_loss.append(to_scalar([loss_recons, loss_vq]))
 
         if (batch_idx + 1) % 100 == 0:
-            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {} Time: {}'.format(
-                epoch, batch_idx * len(data), len(train_loader.dataset),
+            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),
                 time.time() - start_time
             ))
 
 
 def test():
-    x = Variable(test_data[0][0]).cuda()
+    start_time = time.time()
+    val_loss = []
+    for batch_idx, (data, _) in enumerate(test_loader):
+        x = Variable(data, volatile=True).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())
+        val_loss.append(to_scalar([loss_recons, loss_vq]))
+
+    print('\nValidation Completed!\tLoss: {} Time: {:5.3f}'.format(
+        np.asarray(val_loss).mean(0),
+        time.time() - start_time
+    ))
+    return np.asarray(val_loss).mean(0)
+
+
+def generate_samples():
+    x, _ = test_loader.__iter__().next()
+    x = Variable(x[:32]).cuda()
     x_tilde, _, _ = model(x)
-    x_tilde = (x_tilde+1)/2
-    x = (x+1)/2
+    # x_tilde = (x_tilde + 1)/2
+    # x = (x + 1)/2
 
     x_cat = torch.cat([x, x_tilde], 0)
     images = x_cat.cpu().data
     save_image(images, './sample_cifar.png', nrow=8)
 
 
-for i in range(100):
-    train(i)
-    test()
+BEST_LOSS = 999
+LAST_SAVED = -1
+for epoch in range(1, N_EPOCHS):
+    print("Epoch {}:".format(epoch))
+    train()
+    cur_loss, _ = test()
+
+    if cur_loss <= BEST_LOSS:
+        BEST_LOSS = cur_loss
+        LAST_SAVED = epoch
+        print("Saving model!")
+        torch.save(model.state_dict(), 'best_autoencoder.pt')
+    else:
+        print("Not saving model! Last saved: {}".format(LAST_SAVED))
+
+    generate_samples()

modules.py

@@ -1,5 +1,7 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
 
 
 def to_scalar(arr):
@@ -9,13 +11,19 @@ def to_scalar(arr):
         return arr.cpu().data.tolist()[0]
 
 
+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)
+
+
 class ResBlock(nn.Module):
     def __init__(self, dim):
         super(ResBlock, self).__init__()
         self.block = nn.Sequential(
             nn.ReLU(True),
             nn.Conv2d(dim, dim, 3, 1, 1),
-            nn.BatchNorm2d(dim),
             nn.ReLU(True),
             nn.Conv2d(dim, dim, 1)
         )
@@ -25,39 +33,33 @@ def forward(self, x):
 
 
 class AutoEncoder(nn.Module):
-    def __init__(self):
+    def __init__(self, K=512):
         super(AutoEncoder, self).__init__()
         self.encoder = nn.Sequential(
             nn.Conv2d(3, 256, 4, 2, 1),
-            nn.BatchNorm2d(256),
             nn.ReLU(True),
             nn.Conv2d(256, 256, 4, 2, 1),
-            nn.BatchNorm2d(256),
+            nn.ReLU(True),
             ResBlock(256),
-            nn.BatchNorm2d(256),
             ResBlock(256),
-            nn.BatchNorm2d(256)
         )
 
-        self.embedding = nn.Embedding(512, 256)
-        self.embedding.weight.data.copy_(1./512 * torch.randn(512, 256))
+        self.embedding = nn.Embedding(K, 256)
+        self.embedding.weight.data.copy_(1./K * torch.randn(K, 256))
 
         self.decoder = nn.Sequential(
             ResBlock(256),
-            nn.BatchNorm2d(256),
             ResBlock(256),
-            nn.BatchNorm2d(256),
-            nn.ReLU(True),
             nn.ConvTranspose2d(256, 256, 4, 2, 1),
-            nn.BatchNorm2d(256),
             nn.ReLU(True),
             nn.ConvTranspose2d(256, 3, 4, 2, 1),
-            nn.Tanh()
+            nn.Sigmoid()
         )
 
-    def forward(self, x):
+        self.apply(weights_init)
+
+    def encode(self, x):
         z_e_x = self.encoder(x)
-        B, C, H, W = z_e_x.size()
 
         z_e_x_transp = z_e_x.permute(0, 2, 3, 1)  # (B, H, W, C)
         emb = self.embedding.weight.transpose(0, 1)  # (C, K)
@@ -66,8 +68,78 @@ def forward(self, x):
             2
         ).sum(-2)
         latents = dists.min(-1)[1]
+        return latents, z_e_x
 
-        z_q_x = self.embedding(latents.view(latents.size(0), -1))
-        z_q_x = z_q_x.view(B, H, W, C).permute(0, 3, 1, 2)
+    def decode(self, latents):
+        shp = latents.size() + (-1, )
+        z_q_x = self.embedding(latents.view(latents.size(0), -1))  # (B * H * W, C)
+        z_q_x = z_q_x.view(*shp).permute(0, 3, 1, 2)  # (B, C, H, W)
         x_tilde = self.decoder(z_q_x)
+        return x_tilde, z_q_x
+
+    def forward(self, x):
+        latents, z_e_x = self.encode(x)
+        x_tilde, z_q_x = self.decode(latents)
         return x_tilde, z_e_x, z_q_x
+
+
+class MaskedConv2d(nn.Conv2d):
+    def __init__(self, mask_type, *args, **kwargs):
+        super(MaskedConv2d, self).__init__(*args, **kwargs)
+        assert mask_type in {'A', 'B'}
+        self.register_buffer('mask', self.weight.data.clone())
+        _, _, kH, kW = self.weight.size()
+        self.mask.fill_(1)
+        self.mask[:, :, kH // 2, kW // 2 + (mask_type == 'B'):] = 0
+        self.mask[:, :, kH // 2 + 1:] = 0
+
+    def forward(self, x):
+        self.weight.data *= self.mask
+        return super(MaskedConv2d, self).forward(x)
+
+
+class PixelCNN(nn.Module):
+    def __init__(self, dim=64, n_layers=4):
+        super().__init__()
+        self.dim = 64
+
+        # Create embedding layer to embed input
+        self.embedding = nn.Embedding(256, dim)
+
+        # Building the PixelCNN layer by layer
+        net = []
+
+        # Initial block with Mask-A convolution
+        # Rest with Mask-B convolutions
+        for i in range(n_layers):
+            mask_type = 'A' if i == 0 else 'B'
+            net.extend([
+                MaskedConv2d(mask_type, dim, dim, 7, 1, 3, bias=False),
+                nn.BatchNorm2d(dim),
+                nn.ReLU(True)
+            ])
+
+        # Add the output layer
+        net.append(nn.Conv2d(dim, 256, 1))
+
+        self.net = nn.Sequential(*net)
+
+    def forward(self, x):
+        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)
+        return self.net(x)
+
+    def generate(self, batch_size=64):
+        x = Variable(
+            torch.zeros(64, 8, 8).long()
+        ).cuda()
+
+        for i in range(8):
+            for j in range(8):
+                logits = self.forward(x)
+                probs = F.softmax(logits[:, :, i, j], -1)
+                x.data[:, i, j].copy_(
+                    probs.multinomial(1).squeeze().data
+                )
+        return x