<reponame>zelhar/mg21
# based on examples from https://github.com/pytorch/examples
# https://github.com/L1aoXingyu/pytorch-beginner/https://github.com/L1aoXingyu/pytorch-beginner/
# https://github.com/bfarzin/pytorch_aae/blob/master/main_aae.py
# https://github.com/artemsavkin/aae/blob/master/aae.ipynb
import argparse
import os
import torch
import time
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torchvision import datasets, transforms
from torchvision.utils import save_image
from torch.nn import functional as F
import matplotlib.pyplot as plt
from torchvision.utils import make_grid
import numpy as np
import torch.distributions as D

def save_reconstructs(model, x, epoch):
        with torch.no_grad():
            x = x.to(device)
            sample, _, __ = model(x)
            save_image(x.view(x.shape[0], 1, 28, 28),
                       'results/originals_' + str(epoch) + '.png')
            save_image(sample.view(x.shape[0], 1, 28, 28),
                       'results/reconstructs_' + str(epoch) + '.png')

def save_random_reconstructs(model, nz, epoch):
        with torch.no_grad():
            sample = torch.randn(64, nz).to(device)
            sample = model.decode(sample).cpu()
            save_image(sample.view(64, 1, 28, 28),
                       'results/sample_' + str(epoch) + '.png')

class VAE(nn.Module):
    def __init__(self, nin, nz, nh1, nh2, nh3, nh4):
        super(VAE, self).__init__()

        self.nin = nin
        self.nz = nz

        self.encoder = nn.Sequential(
                nn.Linear(nin, nh1), 
                nn.ReLU(),
                nn.Linear(nh1, nh2),
                nn.ReLU(),
                )

        self.decoder = nn.Sequential(
                nn.Linear(nz, nh3),
                nn.ReLU(),
                nn.Linear(nh2, nh4),
                nn.ReLU(),
                nn.Linear(nh4, nin),
                nn.Sigmoid()
                )

        self.mumap = nn.Linear(nh2, nz)
        self.logvarmap = nn.Linear(nh2, nz)

    def encode(self, x):
        h = self.encoder(x)
        mu = self.mumap(h)
        logvar = self.logvarmap(h)
        return mu, logvar
        #h1 = F.relu(self.fc1(x))
        #return self.fc21(h1), self.fc22(h1)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std

    def decode(self, z):
        x = self.decoder(z)
        return x
        #h3 = F.relu(self.fc3(z))
        #return torch.sigmoid(self.fc4(h3))

    def forward(self, x):
        x = x.view(-1, self.nin)
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

# parameters
nin = 28*28
nz = 20
batchSize = 256
epochs = 100

device = "cuda" if torch.cuda.is_available() else "cpu"



train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "../data", train=True, download=True, transform=transforms.ToTensor()
    ),
    batch_size=batchSize,
    shuffle=True,
)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST("../data", train=False, transform=transforms.ToTensor()),
    batch_size=batchSize,
    shuffle=True,
)

bce = nn.BCELoss(reduction="sum")
kld = lambda mu, logvar : -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

model = VAE(nin, nz, 2*1024, 2*512, 2*512, 2*1024).to(device)
optimizer = optim.Adam(model.parameters(), lr=3e-4)


for epoch in range(epochs):
    for idx, (data, _) in enumerate(train_loader):
        batch_size = data.shape[0]
        x = data.view(-1,nin).to(device)
        model.train()
        model.requires_grad_(True)
        optimizer.zero_grad()
        recon, mu, logvar = model(x)
        loss_recon = bce(recon, x)
        loss_kld = kld(mu, logvar)
        #loss_recon.backward()
        #loss_kld.backward()
        loss = loss_kld + loss_recon
        loss.backward()
        optimizer.step()
        if idx % 100 == 0:
            print("losses:\n",
                    "reconstruction loss:", loss_recon.item(),
                    "kld:", loss_kld.item()
                    )


save_random_reconstructs(model, nz, 670)

xs, _ = iter(train_loader).next()

save_reconstructs(model, xs.to(device), 670)

torch.save(model, "results/vae_mnist.pth")

######################################################

mix = D.Categorical(torch.rand(5,)) 
comp = D.Independent(D.Bernoulli(torch.rand(5,2)), 1)
bmm = D.MixtureSameFamily(mix, comp)

comp.sample()
bmm.sample()
bmm.sample((15,))


########################################################
# suppose we are trying to learn n-dimensional data that is generated by
# k n-dimensional independent distributions (mixture model)
# data dimensions:
nin = 50
nk = 2 # number of indipendent distributions in the mixture
nz = 10 # the dimension of out latent space should be >= nk, I think...
batchSize = 528

#weights = torch.tensor([1,2,3,2,1.0])
weights = torch.tensor([1,2])
mix = D.Categorical(weights)
mix.batch_shape, mix.event_shape
mix.sample().shape

#test = D.Normal(torch.rand(3), torch.rand(3))
#comp = D.Independent(D.Bernoulli(torch.rand(5, nin)), 1)
comp = D.Independent(D.Bernoulli(torch.rand(2, nin)), 1)


comp.batch_shape, comp.event_shape

bmm = D.MixtureSameFamily(mix, comp)
bmm.batch_shape, bmm.event_shape

bmm.sample().shape
bmm.sample((10,)).shape



model = VAE(nin, nz, 2024, 912, 912, 2024).to(device)
optimizer = optim.Adam(model.parameters(), lr=3e-4)

TINY = 1e-12

numbatches = 10000
for batch in range(numbatches):
    # generate a batch-some of data
    x = bmm.sample((batchSize,)).to(device)
    x = TINY/2 + (1 - TINY)*x 
    batch_size = x.shape[0]
    # x = data.view(-1,nin).to(device)
    model.train()
    model.requires_grad_(True)
    optimizer.zero_grad()
    recon, mu, logvar = model(x)
    loss_recon = bce(recon + TINY, x)
    loss_kld = kld(mu + TINY, logvar)
    loss = loss_kld + loss_recon
    loss.backward()
    optimizer.step()
    if batch % 100 == 0:
        print(
            "losses:\n",
            "reconstruction loss:",
            loss_recon.item(),
            "kld:",
            loss_kld.item(),
        )



## another test
nin = 5
nk = 2 # number of indipendent distributions in the mixture
nz = 2 # the dimension of out latent space should be >= nk, I think...
batchSize = 128

#weights = torch.tensor([1,2,3,2,1.0])
weights = torch.tensor([1,2])
mix = D.Categorical(weights)
mix.batch_shape, mix.event_shape
mix.sample().shape
compGauss = D.Independent(
        D.Normal(torch.rand(nk, nin), torch.rand(nk, nin)), 1)

gmm = D.MixtureSameFamily(mix, compGauss)
gmm.batch_shape, gmm.event_shape
gmm.sample().shape
gmm.sample((10,)).shape


model = VAE(nin, nz, 2024, 912, 912, 2024).to(device)
optimizer = optim.Adam(model.parameters(), lr=3e-4)
TINY = 1e-12

mse = nn.MSELoss()


numbatches = 30000
for batch in range(numbatches):
    # generate a batch-some of data
    x = gmm.sample((batchSize,)).to(device)
    #x = TINY/2 + (1 - TINY)*x 
    batch_size = x.shape[0]
    # x = data.view(-1,nin).to(device)
    model.train()
    model.requires_grad_(True)
    optimizer.zero_grad()
    recon, mu, logvar = model(x)
    #loss_recon = bce(recon + TINY, x)
    loss_recon = mse(recon + TINY, x)
    loss_kld = kld(mu + TINY, logvar)
    loss = loss_kld + loss_recon
    loss.backward()
    optimizer.step()
    if batch % 100 == 0:
        print(
            "losses:\n",
            "reconstruction loss:",
            loss_recon.item(),
            "kld:",
            loss_kld.item(),
        )




test_real = gmm.sample((1000,)).to(device)
test_fake = torch.randn((1000,5)).to(device)
test_real.shape
test_fake.shape

recon_real, mu_real, logvar_real = model(test_real)
recon_fake, mu_fake, logvar_fake = model(test_fake)

mse(recon_real, test_real)
mse(recon_fake, test_fake)


# save model weights (needs to create model instance to load the weights)
model.cpu()
torch.save(model.state_dict(), "./results/vae_mix_2_5_gauss_weights.pth")
# load model weighte:
model2 = VAE(nin, nz, 2024, 912, 912, 2024).to(device)
model2.load_state_dict(torch.load("./results/vae_mix_2_5_gauss_weights.pth"))
# save the model itself with weights
torch.save(model, "./results/vae_mix_2_5_gauss.pth")
# load model
model3 = torch.load("results/vae_mix_2_5_gauss.pth")












































































