vinf_mnist_unorganized

import torch
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import torch.nn as nn
import torch as t
import torch.nn.init as init
import torch.nn.functional as F
import matplotlib.pyplot as plt
import math
import os 
import numpy as np

#torch.cuda.set_device(4)


device =t.device('cpu')


im_dim = 1
n_classes = 10
imagesize = 32
batchsize = 64

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        '.', train=True,transform = transforms.ToTensor()
    ),
    batch_size=batchsize,
    shuffle=True
)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        '.', train=False, transform = transforms.ToTensor()
    ),
    batch_size=batchsize,
    shuffle=False,
)





D = 28 * 28
K = 16
niters = 50000
data = 'swissroll'
sample = 1
LOW = -4
HIGH = 4
beta = 0.01
alpha = 0.1
hidden_layer1 = 100
hidden_layer2 = 400
hidden_layer3 = 100

latent_dim = 20
lr_decay = 0.99999

# a function for visualizing batch images
def VisConcatImg(batch_images,ax, title):
    batch_size = np.shape(batch_images)[0]
    sqrt_size = int(batch_size ** 0.5)
    batch_images = batch_images.reshape(batch_size, 28, 28)
    row_concatenated = [np.concatenate(batch_images[i*sqrt_size : (i+1)*sqrt_size], axis=1) for i in range(sqrt_size)]
    concatenated = np.concatenate(row_concatenated, axis=0)
    ax.imshow(concatenated, cmap='gray')
    ax.axis('off')
    ax.set_title(title)


def plt_samples(samples, ax, npts=100, title="$x ~ p(x)$"):
    ax.hist2d(samples[:, 0], samples[:, 1], range=[[LOW, HIGH], [LOW, HIGH]], bins=npts, cmap='inferno')
    ax.invert_yaxis()
    ax.get_xaxis().set_ticks([])
    ax.get_yaxis().set_ticks([])
    ax.set_title(title)
    
    
class planar_flow(nn.Module):
    def __init__(self, D, latent_dim):
        super(planar_flow, self).__init__()
        self.D = D
        self.latent_dim = latent_dim
        self.u = nn.Linear(D,latent_dim)
        self.w = nn.Linear(D,latent_dim)
        self.b = nn.Linear(D,1)
        self.h = nn.Tanh()
        #self.init_parameter()
    def forward(self, z ,x):
        u = self.u(x)
        w = self.w(x)
        b = self.b(x)
        linear = self.h(t.sum(z*w, dim = 1) + b.squeeze())
        self.det = t.log(t.abs(1 + t.sum(w*u) * (1 - linear**2)) + 10**(-4))
        return z + (linear[:,None] * u).squeeze()
    
#     def dett(self,z):
#         linear = self.h(F.linear(z, self.w,self.b))
#         return t.log(t.abs(1 + t.sum(self.w*self.u) * (1 - linear**2)) + 10**(-4))
    
#     def init_parameter(self):
#         init.uniform_(self.w,-0.01,0.01)
#         init.uniform_(self.u,-0.01,0.01)
#         init.uniform_(self.b,-0.01,0.01)

inference_net = nn.Sequential(
    nn.Linear(D,hidden_layer1 ),
    nn.ReLU(),
    nn.Linear(hidden_layer1 ,hidden_layer2),
    nn.ReLU(),
    nn.Linear(hidden_layer2 ,hidden_layer3),
    nn.ReLU()
).to(device)

inference_net_mu = nn.Sequential( 
    nn.Linear(hidden_layer3 ,latent_dim)
).to(device)

inference_net_lnsig = nn.Sequential(
        nn.Linear(hidden_layer3 ,latent_dim)
).to(device)

nnet = []
for i in range(K):
    nnet.append(planar_flow(D,latent_dim))
flow_net = nn.Sequential(*nnet).to(device)

generate_net = nn.Sequential( 
    nn.Linear(latent_dim,hidden_layer1 ),
    nn.ReLU(),
    nn.Linear(hidden_layer1 ,hidden_layer2),
    nn.ReLU(),
    nn.Linear(hidden_layer2 ,hidden_layer3),
    nn.ReLU(),
    nn.Linear(hidden_layer3 ,D),
    nn.Sigmoid()
).to(device)



model = nn.Sequential(
inference_net,
inference_net_mu,
inference_net_lnsig,
flow_net,
generate_net
)

#optimizer = t.optim.Adam(model.parameters(),lr=1e-4)
# optimizer = t.optim.Adam([{'params': inference_net_mu.parameters()},
#                 {'params': inference_net_lnsig.parameters()},
#                         {'params': flow_net.parameters()},
#                         {'params': generate_net_mu.parameters()},
#                         {'params': generate_net_lnsig.parameters()}])
optimizer = t.optim.RMSprop(model.parameters(), lr=0.0001,momentum = 0.9)
scheduler = t.optim.lr_scheduler.ExponentialLR(optimizer, lr_decay)
losss_train = []
losss_test = []
temp = 0
for itr in range(niters):
    for i, (x, y) in enumerate(train_loader):
        beta = min(1, 0.01 + temp / 10000)
        scheduler.step()
        optimizer.zero_grad()

        # load data
        x = x.reshape(batchsize,-1).to(device)
        if x.size()[1] != D:
            continue
        epsilon = t.randn(batchsize, latent_dim).to(device)
        inference_x = inference_net(x)
        mu = inference_net_mu(inference_x)
        lnsig = inference_net_lnsig(inference_x)
        sig = t.exp(lnsig)
        z_0 = t.sqrt(sig) * epsilon + mu

        #z_k = flow_net(z_0)
        det = 0
        for net in flow_net:
            z_k = net(z_0,x)
            det = det + net.det
            z_0 = z_k
        
        binary = generate_net(z_0)
          
        l2_reg = None
        for W in model.parameters():
            if l2_reg is None:
                l2_reg = W.norm(2)
            else:
                l2_reg = l2_reg + W.norm(2)
        loss1 = - t.sum(lnsig, dim = 1).mean() / 2
        loss2 = t.sum(z_0 ** 2, dim = 1).mean() / 2
        loss3 = F.binary_cross_entropy(binary,x,reduction  = 'sum') / batchsize
        loss4 = t.mean(det)

        loss = loss1 + beta * (loss2 + loss3) - loss4 
        loss.backward()
        optimizer.step()
        temp = temp + 1

        print('Iter{} |batch {}| Loss {})'.format(itr, i,loss) )
        
        if i % 50 == 0:
            losss_train.append(loss)
            sample_size = 5000
            fig = plt.figure(figsize = (12,4))
            ax1 = plt.subplot(1,3,1)
            ax2 = plt.subplot(1,3,2)
            ax3 = plt.subplot(1,3,3)
                
            with t.no_grad():        
                zz = t.randn(batchsize,latent_dim).to(device)

                gen = generate_net(zz)
  
                VisConcatImg(gen.to('cpu').numpy(), ax3,'generate')
    
                for i, (x, y) in enumerate(train_loader):
                    if i == 1:
                        x = x.reshape(batchsize,-1).to(device)
                        if x.size()[1] != D:
                            continue
                        epsilon = t.randn(batchsize, latent_dim).to(device)
                        inference_x = inference_net(x)
                        mu = inference_net_mu(inference_x)
                        lnsig = inference_net_lnsig(inference_x)
                        sig = t.exp(lnsig)
                        z_0 = t.sqrt(sig) * epsilon + mu

                        for net in flow_net:
                            z_k = net(z_0,x)
                            det = det + net.det
                            z_0 = z_k

                        binary = generate_net(z_0)
                        xxx = binary
                        
                        loss1 = - t.sum(lnsig, dim = 1).mean() / 2
                        loss2 = t.sum(z_0 ** 2, dim = 1).mean() / 2
                        loss3 = F.binary_cross_entropy(binary,x,reduction  = 'sum') / batchsize
                        loss4 = t.mean(det)

                        loss = loss1 + loss2 + loss3 - loss4
                        losss_test.append(loss)
                        
                        VisConcatImg(xxx.to('cpu').numpy(), ax2,'reconstruct')
                        VisConcatImg(x.to('cpu').numpy(), ax1,'raw')
                plt.show()