VAE pytorch

import torch import torch.nn as nn from torch.autograd import Variable import torch.optim as optim from torchvision import transforms,datasets import torch.nn.functional as F import os import scipy import numpy as np from scipy import misc import math batch_size = 64 latent_vector = 32 intermediate_vector = 256 num_class = 10 det = 1e-10 lamb = 2.5 train_data = datasets.MNIST(root='./data/', train=True, transform=transforms.ToTensor(), download= True) test_data = datasets.MNIST(root= './data/', train=True, transform=transforms.ToTensor(), download= True) train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size= batch_size, shuffle= True) test_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size= batch_size, shuffle=True) class VAE(nn.Module): def __init__(self): super(VAE, self).__init__() self.encoder = nn.Sequential( nn.Conv2d(1, 32, 3, padding=1, stride=1), nn.LeakyReLU(negative_slope=0.2), nn.Conv2d(32, 32, 3, padding=1, stride=2), nn.LeakyReLU(negative_slope=0.2), nn.Conv2d(32, 64, 3, padding=1, stride=1), nn.LeakyReLU(negative_slope=0.2), nn.Conv2d(64, 64, 3, padding=1, stride=2), # 7*7*64 nn.LeakyReLU(negative_slope=0.2), ) self.fc_mu = nn.Linear(7*7*64, latent_vector) self.fc_logvar = nn.Linear(7*7*64, latent_vector) self.fc = nn.Linear(latent_vector, 7*7*64) self.decoder = nn.Sequential( nn.ConvTranspose2d(64, 64, 4, padding=1, stride=2), nn.LeakyReLU(negative_slope=0.2), nn.ConvTranspose2d(64, 32, 3, padding=1, stride=1), nn.LeakyReLU(negative_slope=0.2), nn.ConvTranspose2d(32, 32, 4, padding=1, stride=2), nn.LeakyReLU(negative_slope=0.2), nn.ConvTranspose2d(32, 1, 3, padding=1, stride=1), nn.Sigmoid() ) def Reparameter(self, mu, logvar): parameter = Variable(torch.randn(mu.size(0), mu.size(1))) return parameter * torch.exp(logvar/2) + mu def forward(self, x): mu = self.encoder(x) logvar = self.encoder(x) mu = self.fc_mu(mu.view(mu.size(0), -1)) logvar = self.fc_logvar(logvar.view(logvar.size(0), -1)) z = self.Reparameter(mu, logvar) vector = self.fc(z).view(z.size(0), 64, 7, 7) return self.decoder(vector), mu, logvar model = VAE() MSE_Loss = nn.MSELoss(size_average=False) #MSE_Loss = F.binary_cross_entropy(reduction='sum') def loss_function(input, output, mu, logvar): Mse = MSE_Loss(input, output) # Mse = 0.5 * torch.mean((input - output).pow(2), 0) KL_loss = 0.5 * torch.sum(-logvar + mu.pow(2) + logvar.exp() - 1) # KL_loss = -0.5 * (temp_logvar - z_se.pow(2)) # # KL_loss = torch.mean(torch.tensordot(torch.unsqueeze(y, 1), KL_loss), 0) # # cat_loss = torch.mean(y * torch.log(y + det), 0) # # return lamb * torch.sum(Mse) + torch.sum(KL_loss) + torch.sum(cat_loss) return Mse + KL_loss def save_image(output, size, path, Color): h, w = output.shape[1], output.shape[2] if Color is True: image = np.zeros((w * size[0], h * size[1], 3)) else: image = np.zeros((w * size[0], h * size[1])) for index, data in enumerate (output): i = index % size[0] j = math.floor(index / size[1]) if Color is True: image[h*j : h*j+h, w*i : w*i+w, :] = data else: image[h*j : h*j+j, w*i : w*i+w] = data scipy.misc.toimage((image*255), cmin=0, cmax=255).save(path) def rescale_image(image): return (image/1.5+0.5)*255 optimizer = optim.SGD(model.parameters(), lr= 0.0001) def train(): for epoch in range(1,10): for i, (data, _) in enumerate (train_loader): tensor_data = Variable(data) output, mu, logvar= model(tensor_data) optimizer.zero_grad() loss = loss_function(tensor_data, output, mu, logvar) loss.backward() optimizer.step() if i % 50 == 0: if not os.path.exists("./image"): os.mkdir("./image") np_output = output.detach().numpy() np_output = np_output.swapaxes(1,2).swapaxes(2,3) save_image(np_output, [8,8], './image/image_{}.png'.format(i), True) print("loss={}".format(loss)) train()