Third commit

Author: ayushtues

.idea/GenZoo.iml

@@ -1,6 +1,8 @@
 import torch
 import torch.utils.data
 from torchvision import datasets, transforms
+import torch.nn as nn
+import torch.nn.functional as F
 
 
 def load_mnist(batch_size):
@@ -10,6 +12,3 @@ def load_mnist(batch_size):
         batch_size=batch_size, shuffle=True)
 
     return train_loader
-
-
-

.idea/dictionaries/ayushtues.xml

@@ -0,0 +1,26 @@
+import torch
+import torch.utils.data
+from torchvision import datasets, transforms
+from torch import nn, optim
+import model as VAE
+import torch.nn.functional as F
+import data_loader as load
+from tensorboardX import SummaryWriter
+import  train.py as train_model
+import matplotlib
+import matplotlib.pyplot as plt
+
+
+trainloader = load.load_mnist(batch_size=60)
+x = trainloader.next()
+plt.imshow(x.numpy()[0], cmap='gray')
+
+model = VAE.make_model()
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+device = ('cuda' if torch.cuda.is_available() else 'cpu')
+epoch = 10
+
+train_model.train(trainloader, epoch, optimizer, model, device)
+
+
+

VAE/data_loader.py

@@ -2,6 +2,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+__name__ = "model.py"
+
 
 class VAE(nn.Module):
     def __init__(self, z_dim=20, keep_prob=0.2):
@@ -14,11 +16,11 @@ def __init__(self, z_dim=20, keep_prob=0.2):
         self.fc3 = nn.Linear(z_dim, 256)
         self.fc4 = nn.Linear(256, 512)
         self.fc5 = nn.Linear(512, 784)
-        self.decode = nn.Sequential(nn.ConvTranspose2d(16, 16, 3, padding=1),
-                                    nn.BatchNorm2d(16),
-                                    nn.ConvTranspose2d(16, 3, 8),
+        self.decode = nn.Sequential(nn.ConvTranspose2d(16, 10, 3, padding=1),
+                                    nn.BatchNorm2d(10),
+                                    nn.ConvTranspose2d(10, 5, 8),
                                     nn.BatchNorm2d(3),
-                                    nn.ConvTranspose2d(3, 1, 15))
+                                    nn.ConvTranspose2d(5, 1, 15))
 
         self.mean = nn.Linear(256, z_dim)
         self.logvar = nn.Linear(256, z_dim)
@@ -37,13 +39,29 @@ def encoder(self, x):
         return mean, logvar
 
     def reparameterize(self, mean, logvar):
-        std = logvar.mul(0.5).exp_()
-        eps = torch.randn_like(std)
-        return mean + eps * std
+        if self.training:
+            std = logvar.mul(0.5).exp_()
+            eps = torch.randn_like(std)
+            return mean + eps * std
+        else:
+            return mean
 
     def decoder(self, z):
         z = self.fc3(z)
         z = self.fc4(z)
         z = self.fc5(z)
         z = z.view([-1, 16, 7, 7])
+        z = self.decode(z)
+        z = F.sigmoid(z)
+        return z
+
+    def forward(self, x):
+        mean, logvar = self.encoder(x)
+        z = self.reparameterize(mean, logvar)
+        x_output = self.decoder(z)
+        return x_output, mean, logvar
+
 
+def make_model():
+    model = VAE()
+    return model

VAE/main.py

@@ -0,0 +1,27 @@
+import torch
+import torchvision
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import nn, optim
+import model as VAE
+
+batch_size = 20
+
+
+def loss_function(x_output, x, mean, logvar):
+    bce = F.binary_cross_entropy(x_output, x.view(-1, 784))
+    kld = -0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp())
+    return (bce + kld) / (784 * batch_size)
+
+
+def train(trainloader, epoch, optimiser, model, device):
+    for i in range(epoch):
+        model.train()
+        model.to(device)
+        for images, _ in trainloader:
+            images = images.to(device)
+            optimiser.zero_grad()
+            x_output, mean, logvar = model(images)
+            loss = loss_function(x_output, images, mean, logvar)
+            loss.backward()
+            optimiser.step()