PMR tutorial on VAEs¶

# Import main libraries
import argparse
import itertools
import random
import os
from tqdm import tqdm
from copy import deepcopy

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
from IPython import display

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim

def seed_everything(seed=20211201):
    """
    Helper to seed random number generators.
    
    Note, however, that even with the seeds fixed, some non-determinism is possible. 
    
    For more details read <https://pytorch.org/docs/stable/notes/randomness.html>.
    """
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    # The below might help to make it more deterministic
    # but will hurt the performance significantly
    # torch.backends.cudnn.benchmark = False
    # torch.backends.cudnn.deterministic = True

img_size=14

train_data = torchvision.datasets.MNIST(
    './data', train=True, download=True,
    transform=torchvision.transforms.Compose([torchvision.transforms.Resize((img_size,img_size)),
                                              torchvision.transforms.ToTensor()]))
test_data = torchvision.datasets.MNIST(
    './data', train=False,
    transform=torchvision.transforms.Compose([torchvision.transforms.Resize((img_size,img_size)),
                                              torchvision.transforms.ToTensor()]))

data_dims=train_data[0][0].cpu().numpy().shape[1:]

def plot_grid_of_mnist_samples(samples, grid=(5, 5), figsize=(8,8)):
    """A helper furnction for plotting MNIST samples"""
    fig, axes = plt.subplots(grid[0], grid[1], figsize=figsize)
    axes = axes.flatten()
    for idx, ax in enumerate(axes):
        ax.imshow(samples[idx], cmap='gray')

        ax.tick_params(left=False, labelleft=False,
                       bottom=False, labelbottom=False)

    fig.tight_layout()
    return fig

def evaluate_logprob_continuous_bernoulli(X, *, logits):
    """
    Evaluates log-probability of the continuous Bernoulli distribution

    Args:
        X (Tensor):      data, a batch of shape (B, D)
        logits (Tensor): parameters of the continuous Bernoulli,
                         a batch of shape (B, D) 

    Returns:
        logpx (Tensor): log-probabilities of the inputs X, a batch of shape (B,)
    """
    cb = torch.distributions.ContinuousBernoulli(logits=logits)
    return cb.log_prob(X).sum(dim=-1)

def evaluate_logprob_diagonal_gaussian(Z, *, mean, logvar):
    """
    Evaluates log-probability of the diagonal Gaussian distribution

    Args:
        Z (Tensor):      latent vectors, a batch of shape (*, B, H)
        mean (Tensor):   mean of diagonal Gaussian, a batch of shape (*, B, H)
        logvar (Tensor): log-variance of diagonal Gaussian, a batch of shape (*, B, H)

    Returns:
        logqz (Tensor): log-probabilities of the inputs Z, a batch of shape (*, B)
                        where `*` corresponds to any additional dimensions of the input arguments,
                        for example a dimension representing the samples used to approximate 
                        the expectation in the ELBO
    """
    gauss = torch.distributions.Normal(loc=mean, scale=torch.exp(0.5*logvar))
    return gauss.log_prob(Z).sum(dim=-1)

def compute_kld_with_standard_gaussian(q_mean, q_logvar):
    """
    Computes KL(q(z|x)||p(z)) between the variational diagonal 
    Gaussian distribution q and standard Gaussian prior p
    
    Args:
        q_mean (Tensor):   mean of diagonal Gaussian q, a batch of shape (B, H)
        q_logvar (Tensor): log-variance of diagonal Gaussian q, a batch of shape (B, H)
        
    Returns:
        kld (Tensor): KL divergences between q and p, a batch of shape (B,)
    """
    ### IMPLEMENT YOUR CODE HERE
    
    # q_var = ...
    
    # kld = ...
    
    q_var = q_logvar.exp()
    
    kld = -0.5 * (1 + q_logvar - q_mean**2 - q_var).sum(dim=-1)
    
    ###
    return kld

def sample_gaussian_with_reparametrisation(mean, logvar, *, num_samples=1):
    """
    Samples the Gaussian distribution using the reparametrisation trick
    
    Args:
        mean (Tensor):     mean of diagonal Gaussian q, a batch of shape (B, K)
        logvar (Tensor):   log-variance of diagonal Gaussian q, a batch of shape (B, K)
        num_samples (int): The number of samples (M) to approximate the expectation in the ELBO
    
    Returns:
        Z (Tensor):        Samples Z from the diagonal Gaussian q, a batch of shape (num_samples, B, K)
    """
    ### IMPLEMENT YOUR CODE HERE
    
    # std = ...
    # eps = ...
    
    # Z = ...
    
    std = torch.exp(0.5*logvar)
    eps = torch.randn(torch.Size([num_samples]) + std.shape,
                      dtype=std.dtype, device=std.device)
    
    Z = mean + eps*std
    ###
    
    # Making sure that the variational samples are on the same device as the input argument (i.e. CPU or GPU)
    return Z.to(device=std.device)

def elbo_with_pathwise_gradients(X, *, encoder, decoder, num_samples):
    """
    Evaluate the ELBO for each data-point in X
    
    Args:
        X (Tensor):                 data, a batch of shape (B, D)
        decoder (torch.nn.Module):  is a neural network that provides the parameters
                                    of the generative model p(x|z; θ)
        encoder (torch.nn.Module):  is a neural network that provides the parameters 
                                    of the variational distribution q(z|x; φ).
        num_samples (int):          the number of samples used in the Monte Carlo averaging

    Returns:
        elbos (Tensor):             the ELBOs for each input X, a tensor of shape (B,)
    """
    # Evaluate the encoder network to obtain the parameters of the 
    # variational distribution
    mean, logvar = encoder(X)  # (B, H), (B, H)
    
    # Sample the latents using the reparametrisation trick
    Z = sample_gaussian_with_reparametrisation(
            mean, logvar, num_samples=num_samples)  # (M, B, H)

    # Evaluate the decoder network to obtain the parameters of the 
    # generative model p(x|z)
    logits = decoder(Z)  # (M, B, D)
    
    ### IMPLEMENT YOUR CODE BELOW
    
    # KLD( q(z|x) || p(z) )
    # kld = ...  # (B,)

    # ~E_{q(z|x)}[ p(x | z) ]
    # neg_cross_entropy = ...  # (B,)

    # ELBO for each data-point
    # elbos = ...  # (B,)

    # KLD( q(z|x) || p(z) )
    kld = compute_kld_with_standard_gaussian(mean, logvar)  # (B,)

    # ~E_{q(z|x)}[ p(x | z) ]
    neg_cross_entropy = evaluate_logprob_continuous_bernoulli(X, logits=logits).mean(dim=0)  # (B,)
    
    # ELBO for each data-point
    elbos = neg_cross_entropy - kld  # (B,)
    ###
    
    return elbos

def sample_gaussian_without_reparametrisation(mean, logvar, *, num_samples=1):
    """
    Samples the Gaussian distribution without attaching then to the computation graph
    
    Args:
        mean (Tensor):   mean of diagonal Gaussian q, a batch of shape (B, K)
        logvar (Tensor): log-variance of diagonal Gaussian q, a batch of shape (B, K)
        num_samples (int): The number of samples (M) to approximate the expectation in the ELBO
    
    Returns:
        Z (Tensor):      Samples Z from the diagonal Gaussian q, a batch of shape (num_samples, B, K)
    """
    return sample_gaussian_with_reparametrisation(mean, logvar, num_samples=num_samples).detach()

def elbo_with_score_function_gradients(X, *, encoder, decoder, num_samples):
    """
    Evaluate the surrogate ELBO for each data-point in X. The surrogate
    elbo uses score function gradient estimator to estimate the gradients
    of the variational distribution.
    
    Args:
        X (Tensor):                 data, a batch of shape (B, D)
        decoder (torch.nn.Module):  is a neural network that provides the parameters
                                    of the generative model p(x|z; θ)
        encoder (torch.nn.Module):  is a neural network that provides the parameters 
                                    of the variational distribution q(z|x; φ).
        num_samples (int):          the number of samples used in the Monte Carlo averaging
    
    Returns:
        elbos (Tensor):             the ELBOs for each input X, a tensor of shape (B,)
    """
    # Evaluate the encoder network to obtain the parameters of the 
    # variational distribution
    mean, logvar = encoder(X)  # (B, H), (B, H)
    
    # Sample the latents _without_ the reparametrisation trick
    Z = sample_gaussian_without_reparametrisation(
            mean, logvar, num_samples=num_samples)  # (M, B, H)

    # Evaluate the decoder network to obtain the parameters of the 
    # generative model p(x|z)
    logits = decoder(Z)  # (M, B, D)
    
    ### IMPLEMENT YOUR CODE BELOW
    
    #
    # ELBO
    #
    
    # KLD( q(z|x) || p(z) )
    # kld = ...

    # log p(x | z)
    # log_px_given_z = ...
    
    # ~E_{q(z|x)}[ p(x | z) ]
    # neg_cross_entropy = ...
    
    # Per-data-point ELBO that does not estimate the gradients of the expectation parameters
    # elbos = ...
    
    #
    # Score function surrogate
    #

    # log q(z | x)
    # log_qz_given_x = ...

    # Surrogate loss for score function gradient estimator of the expectation parameters
    # surrogate_loss = ...

    # ELBO with score function gradients
    # surrogate_elbos = ...
    
    #
    # ELBO
    #

    # KLD( q(z|x) || p(z) )
    kld = compute_kld_with_standard_gaussian(mean, logvar)

    # log p(x | z)
    log_px_given_z = evaluate_logprob_continuous_bernoulli(X, logits=logits)
    
    # ~E_{q(z|x)}[ p(x | z) ]
    neg_cross_entropy = log_px_given_z.mean(dim=0)
    
    # Per data-point ELBO, note that this ELBO will not estimate the gradients of the expectation parameters!
    elbos = neg_cross_entropy - kld
    
    #
    # Score function surrogate
    #

    # log q(z | x)
    log_qz_given_x = evaluate_logprob_diagonal_gaussian(Z, mean=mean, logvar=logvar)

    # Surrogate loss for score function gradient estimator
    surrogate_loss = (log_qz_given_x*log_px_given_z.detach()).mean(dim=0)

    # ELBO with score function gradients
    surrogate_elbos = elbos + surrogate_loss - surrogate_loss.detach()
    
    ###
    return surrogate_elbos

class Encoder(nn.Module):
    """
    Encoder or inference network that predicts the parameters of the variational distribution q(z|x).
    """
    def __init__(self, hparams):
        super().__init__()
        self.hparams = hparams
        enc_layer_dims = self.hparams.enc_layer_dims

        # Create all layers except last
        layers = []
        for i in range(len(enc_layer_dims)-2):
            layers.append(nn.Linear(enc_layer_dims[i],
                                    enc_layer_dims[i+1]))
            # Use a non-linearity
            layers.append(nn.ReLU())

        self.model = nn.Sequential(*layers)

        # Create separate final layers for each parameter (mean and log-variance)
        # We use log-variance to unconstrain the optimisation of the positive-only variance parameters
        self.mean = nn.Linear(enc_layer_dims[-2],
                              enc_layer_dims[-1])
        self.logvar = nn.Linear(enc_layer_dims[-2],
                                enc_layer_dims[-1])

    @staticmethod
    def add_model_args(parser):
        """Here we define the arguments for our encoder model."""
        parser.add_argument('--enc_layer_dims', type=int, nargs='+',
                            help='Encoder layer dimensions.')
        return parser
    
    def forward(self, X):
        """
        Predicts the parameters of the variational distribution
        
        Args:
            X (Tensor):      data, a batch of shape (B, D)
            
        Returns:
            mean (Tensor):   means of the variational distributions, shape (B, K)
            logvar (Tensor): log-variances of the diagonal Gaussian variational distribution, shape (B, K)
        """
        features = self.model(X)
        mean = self.mean(features)
        logvar = self.logvar(features)
        
        return mean, logvar

    def sample_with_reparametrisation(self, mean, logvar, *, num_samples=1):
        # Reuse the implemented code
        return sample_gaussian_with_reparametrisation(mean, logvar, num_samples=num_samples)

    def sample_without_reparametrisation(self, mean, logvar, *, num_samples=1):
        # Reuse the implemented code
        return sample_gaussian_without_reparametrisation(mean, logvar, num_samples=num_samples)
    
    def log_prob(self, mean, logvar, Z):
        """
        Evaluates the log_probability of Z given the parameters of the diagonal Gaussian
        
        Args:
            mean (Tensor):   means of the variational distributions, shape (*, B, K)
            logvar (Tensor): log-variances of the diagonal Gaussian variational distribution, shape (*, B, K)
            Z (Tensor):      latent vectors, shape (*, B, K)
        
        Returns:
            logqz (Tensor):  log-probability of Z, a batch of shape (*, B)
        """
        # Reuse the implemented code
        return evaluate_logprob_diagonal_gaussian(Z, mean=mean, logvar=logvar)

class Decoder(nn.Module):
    """
    Decoder or generative network that computes the parameters of the likelihood.
    """
    def __init__(self, hparams):
        super().__init__()
        self.hparams = hparams
        dec_layer_dims = self.hparams.dec_layer_dims
        
        # Create all layers except last
        layers = []
        for i in range(len(dec_layer_dims)-2):
            layers.append(nn.Linear(dec_layer_dims[i],
                                    dec_layer_dims[i+1]))
            # Add non-linearity
            layers.append(nn.ReLU())

        self.model = nn.Sequential(*layers)
        # Create final layers that predicts the parameters of the continuous Bernoulli
        self.logits = nn.Linear(dec_layer_dims[-2],
                                dec_layer_dims[-1])
        
    @staticmethod
    def add_model_args(parser):
        """Here we define the arguments for our decoder model."""
        parser.add_argument('--dec_layer_dims', type=int, nargs='+',
                            help='Decoder layer dimensions.')
        return parser

    def forward(self, Z):
        """
        Computes the parameters of the generative distribution p(x | z)
        
        Args:
            Z (Tensor):  latent vectors, a batch of shape (M, B, K)
            
        Returns:
            logits (Tensor):   parameters of the continuous Bernoulli, shape (M, B, D)
        """
        features = self.model(Z)
        logits = self.logits(features)

        return logits

    def log_prob(self, logits, X):
        """
        Evaluates the log_probability of X given the parameters of the continuous Bernoulli
        
        Args:
            logits (Tensor): parameters of the continuous Bernoulli, shape (*, B, D)
            X (Tensor):      data, shape (*, B, D)
        
        Returns:
            logpx (Tensor):  log-probability of X, a batch of shape (*, B)
        """
        # Reuse the implemented code
        return evaluate_logprob_continuous_bernoulli(X, logits=logits)
    
    # Some extra methods for analysis
    
    def sample(self, logits, *, num_samples=1):
        """
        Samples the continuous Bernoulli
        
        Args:
            logits (Tensor):   parameters of the continuous Bernoulli, shape (*, B, D)
            num_samples (int): number of samples
        
        Returns:
            X (Tensor):  samples from the distribution, shape (num_samples, *, B, D)
        """
        cb = torch.distributions.ContinuousBernoulli(logits=logits)
        return cb.sample((num_samples,))

    def mean(self, logits):
        """
        Returns the mean of the continuous Bernoulli
        
        Args:
            logits (Tensor):   parameters of the continuous Bernoulli, shape (*, B, D)
        
        Returns:
            mean (Tensor):  means of the continuous Bernoulli, shape (*, B, D)
        """
        cb = torch.distributions.ContinuousBernoulli(logits=logits)
        return cb.mean
    
    def param_p(self, logits):
        """
        Returns the success probability p of the continuous Bernoulli
        
        Args:
            logits (Tensor):   parameters of the continuous Bernoulli, shape (*, B, D)
        
        Returns:
            p (Tensor):  success probability p of the continuous Bernoulli, shape (*, B, D)
        """
        cb = torch.distributions.ContinuousBernoulli(logits=logits)
        return cb.probs

class VAE(nn.Module):
    """
    A wrapper for the VAE model
    """
    def __init__(self, hparams):
        super().__init__()
        self.hparams = hparams
        
        # Use the encoder and decoder implemented above
        self.encoder = Encoder(hparams)
        self.decoder = Decoder(hparams)
        
    @staticmethod
    def add_model_args(parser):
        """Here we define the arguments for our decoder model."""
        parser = Encoder.add_model_args(parser)
        parser = Decoder.add_model_args(parser)
        
        parser.add_argument('--gradient_type', type=str,
                            choices=['pathwise', 'score'],
                            help='Variational model gradient estimation method.')
        parser.add_argument('--num_variational_samples', 
                            type=int, default=1,
                            help=('The number of samples from the variational '
                                  'distribution to approximate the expectation.'))
        
        return parser

    def forward(self, X):
        """
        Computes the variational ELBO
        
        Args:
            X (Tensor):  data, a batch of shape (B, K)
            
        Returns:
            elbos (Tensor): per data-point elbos, shape (B, D)
        """
        if self.hparams.gradient_type == 'pathwise':
            return self.elbo_with_pathwise_gradients(X)
        elif self.hparams.gradient_type == 'score':
            return self.elbo_with_score_function_gradients(X)
        else:
            raise ValueError(f'gradient_type={self.hparams.gradient_type} is invalid')
            
    def elbo_with_pathwise_gradients(self, X):
        # Reuse the implemented code
        return elbo_with_pathwise_gradients(X, encoder=self.encoder, decoder=self.decoder, 
                                            num_samples=self.hparams.num_variational_samples)
        
    def elbo_with_score_function_gradients(self, X):
        # Reuse the implemented code
        return elbo_with_score_function_gradients(X, encoder=self.encoder, decoder=self.decoder, 
                                                  num_samples=self.hparams.num_variational_samples)

def per_datapoint_elbo_to_avgelbo_and_loss(elbos):
    # Compute the average ELBO over the mini-batch
    elbo = elbos.mean(0)
    # We want to _maximise_ the ELBO, but the SGD implementations 
    # do minimisation by default, hence we multiply the ELBO by -1.
    loss = -elbo
    
    return elbo, loss

def create_dataloader(data, args):
    kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}

    return torch.utils.data.DataLoader(
        data, batch_size=args.batch_size, shuffle=True, **kwargs)

def fit_model(model, optimizer, train_data, args, *, test_data=None):
    # We will plot the learning curves during training
    fig, ax = plt.subplots(1, 1, figsize=(8, 4))
    
    # Create data loaders
    train_loader = create_dataloader(train_data, args)
    if test_data is not None:
        test_loader = create_dataloader(test_data, args)

    train_epochs = []
    train_elbos = []
    train_avg_epochs = []
    train_avg_elbos = []
    test_avg_epochs = []
    test_avg_elbos = []
    
    # We will use these to track the best performing model on test data
    best_avg_test_elbo = float('-inf')
    best_epoch = None
    best_model_state = None
    best_optim_state = None

    pbar = tqdm(range(1, args.epochs + 1))
    for epoch in pbar:
        # Train
        model.train()
        epoch_train_elbos = []
        # We don't use labels hence discard them with a _
        for batch_idx, (mbatch, _) in enumerate(train_loader):
            mbatch = mbatch.to(args.device)
            # Flatten the images
            mbatch = mbatch.view([-1] + [mbatch.shape[-2]*mbatch.shape[-1]])
            # Reset gradient computations in the computation graph
            optimizer.zero_grad()

            # Compute the loss for the mini-batch
            elbo, loss = per_datapoint_elbo_to_avgelbo_and_loss(model(mbatch))

            # Compute the gradients using backpropagation
            loss.backward()
            # Perform an SGD update
            optimizer.step()

            epoch_train_elbos += [elbo.detach().item()]
            pbar.set_description((f'Train Epoch: {epoch} [{batch_idx * len(mbatch)}/{len(train_loader.dataset)}'
                                  f'({100. * batch_idx / len(train_loader):.0f}%)] ELBO: {elbo:.6f}'))

        # Test
        if test_data is not None:
            with torch.inference_mode():
                model.eval()
                epoch_test_elbos = []
                for batch_idx, (mbatch, _) in enumerate(test_loader):
                    mbatch = mbatch.to(args.device)
                    # Flatten the images
                    mbatch = mbatch.view([-1] + [mbatch.shape[-2]*mbatch.shape[-1]])

                    # Compute the loss for the test mini-batch
                    elbo, loss = per_datapoint_elbo_to_avgelbo_and_loss(model(mbatch))

                    epoch_test_elbos += [elbo.detach().item()]
                    pbar.set_description((f'Test Epoch: {epoch} [{batch_idx * len(mbatch)}/{len(test_loader.dataset)} '
                                          f'({100. * batch_idx / len(test_loader):.0f}%)] ELBO: {elbo:.6f}'))

        # Store epoch summary in list
        train_avg_epochs += [epoch]
        train_avg_elbos += [np.mean(epoch_train_elbos)]
        train_epochs += np.linspace(epoch-1, epoch, len(epoch_train_elbos)).tolist()
        train_elbos += epoch_train_elbos
        if test_data is not None:
            test_avg_epochs += [epoch]
            epoch_avg_test_elbo = np.mean(epoch_test_elbos)
            test_avg_elbos += [epoch_avg_test_elbo]
            
            # Snapshot best model
            if epoch_avg_test_elbo > best_avg_test_elbo:
                best_avg_test_elbo = epoch_avg_test_elbo
                best_epoch = epoch
                
                best_model_state = deepcopy(model.state_dict())
                best_optim_state = deepcopy(optimizer.state_dict())

        # Update learning curve figure
        ax.clear()
        ax.plot(train_epochs, train_elbos, color='b', alpha=0.5, label='train')
        ax.plot(np.array(train_avg_epochs)-0.5, train_avg_elbos, color='b', label='train (avg)')
        if len(test_avg_elbos) > 0:
            ax.plot(np.array(test_avg_epochs)-0.5, test_avg_elbos, color='r', label='test (avg)')
        ax.grid(True)
        ax.xaxis.set_major_locator(mpl.ticker.MaxNLocator(integer=True))
        ax.legend(loc='lower right')
        ax.set_ylabel('ELBO')
        ax.set_xlabel('Epoch')

        fig_title = f'Epoch: {epoch}, Avg. train ELBO: {np.mean(epoch_train_elbos):.2f}, Avg. test ELBO: {np.mean(epoch_test_elbos):.2f}'
        # If we are tracking best model, then also highlight it on the plot and figure title
        if best_avg_test_elbo != float('-inf'):
            fig_title += f', Best avg. test ELBO: {best_avg_test_elbo:.2f}'
            ax.scatter(best_epoch-0.5, best_avg_test_elbo, marker='*', color='r')
            

        fig.suptitle(fig_title, size=13)
        fig.tight_layout()
        display.clear_output(wait=True)
        if epoch != args.epochs:
            # Force display of the figure (except last epoch, where 
            # jupyter automatically shows the contained figure)
            display.display(fig)
        
    # Reset gradient computations in the computation graph
    optimizer.zero_grad()
    
    if best_model_state is not None and best_epoch != args.epochs:
        print(f'Loading best model state from epoch {best_epoch}.')
        model.load_state_dict(best_model_state)
    if best_optim_state is not None and best_epoch != args.epochs:
        print(f'Loading best optimizer state from epoch {best_epoch}.')
        optimizer.load_state_dict(best_optim_state)
        
    out = {
        'train_avg_epochs': train_avg_epochs,
        'train_avg_elbos': train_avg_elbos,
        'train_epochs': train_epochs,
        'train_elbos': train_elbos,
        'test_avg_epochs': test_avg_epochs,
        'test_avg_elbos': test_avg_elbos
    }
    return model, optimizer, out, fig

def save_state(model, optimizer, file):
    return torch.save({
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            }, file)

def load_state(model, optimizer, file):
    checkpoint = torch.load(file)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    return

# Specify the hyperpameter choices
data_dim = train_data[0][0].numel()
args_path = ['--dec_layer_dims', '2', '100', f'{data_dim}',
             '--enc_layer_dims', f'{data_dim}', '100', '2', 
             '--gradient_type', 'pathwise',
             '--num_variational_samples', '1',
             '--epochs', '50',
             '--learning_rate', '1e-3',
             '--cuda']

args_path = parser.parse_args(args_path)
args_path.cuda = args_path.cuda and torch.cuda.is_available()
args_path.device = torch.device("cuda" if args_path.cuda else "cpu")

# Seed all random number generators for reproducibility of the runs
seed_everything(args_path.seed)

# Initialise the model and the Adam (SGD) optimiser
model_path = VAE(args_path).to(args_path.device)
optimizer_path = optim.Adam(model_path.parameters(), lr=args_path.learning_rate)

# save_state(model_path, optimizer_path, 'VAE_path.pt')

# Alternatively to training a model from scratch you can load a pre-trained model
# load_state(model_path, optimizer_path, 'VAE_path.pt')

# Specify the hyperpameter choices
data_dim = train_data[0][0].numel()
args_score = ['--dec_layer_dims', '2', '100', f'{data_dim}',
              '--enc_layer_dims', f'{data_dim}', '100', '2',
              '--gradient_type', 'score',
              '--num_variational_samples', '1',
              '--epochs', '50',
              '--learning_rate', '1e-4',
              '--cuda']

args_score = parser.parse_args(args_score)
args_score.cuda = args_score.cuda and torch.cuda.is_available()
args_score.device = torch.device("cuda" if args_score.cuda else "cpu")

# Seed all random number generators for reproducibility of the runs
seed_everything(args_score.seed)

# Initialise the model and the Adam (SGD) optimiser
model_score = VAE(args_score).to(args_score.device)
optimizer_score = optim.Adam(model_score.parameters(), lr=args_score.learning_rate)

# save_state(model_score, optimizer_score, 'VAE_score.pt')

# Alternatively to training a model from scratch you can load a pre-trained model
# load_state(model_score, optimizer_score, 'VAE_score.pt')

class VAEwScoreVR(VAE):
    """VAE using score function gradient estimator and variance reduction"""
    
    @staticmethod
    def add_model_args(parser):
        """Here we define the arguments for our encoder model."""
        parser = super(VAEwScoreVR, VAEwScoreVR).add_model_args(parser)
        parser.add_argument('--vr_coeff', type=float,
                            help='Variance reduction coefficient.')
        return parser

    def elbo_with_score_function_gradients(self, X):
        # Evaluate the encoder network to obtain the parameters of the 
        # variational distribution
        mean, logvar = self.encoder(X)  # (B, H), (B, H)

        # Sample the latents _without_ the reparametrisation trick
        Z = sample_gaussian_without_reparametrisation(
                mean, logvar, num_samples=self.hparams.num_variational_samples)  # (M, B, H)

        # Evaluate the decoder network to obtain the parameters of the 
        # generative model p(x|z)
        logits = self.decoder(Z)  # (M, B, D)
        
        #
        # ELBO
        #

        # KLD( q(z|x) || p(z) )
        kld = compute_kld_with_standard_gaussian(mean, logvar)

        # log p(x | z)
        log_px_given_z = evaluate_logprob_continuous_bernoulli(X, logits=logits)

        # ~E_{q(z|x)}[ p(x | z) ]
        neg_cross_entropy = log_px_given_z.mean(dim=0)

        # Per data-point ELBO, note that this ELBO will not estimate the gradients of the expectation parameters!
        elbos = neg_cross_entropy - kld
        
        #
        # Score function surrogate
        #

        # log q(z | x)
        log_qz_given_x = evaluate_logprob_diagonal_gaussian(Z, mean=mean, logvar=logvar)

        # Surrogate loss for score function gradient estimator
        surrogate_loss = (log_qz_given_x*log_px_given_z.detach()).mean(dim=0)

        # ELBO with score function gradients
        surrogate_elbos = elbos + surrogate_loss - surrogate_loss.detach()
        
        #
        # Variance reduction
        # 

        avg_log_qz_given_x = log_qz_given_x.mean(dim=0)
        surrogate_elbos -= self.hparams.vr_coeff*(avg_log_qz_given_x - avg_log_qz_given_x.detach())

        ###
        return surrogate_elbos

# Specify the hyperpameter choices
data_dim = train_data[0][0].numel()
args_score_vr = ['--dec_layer_dims', '2', '100', f'{data_dim}',
                 '--enc_layer_dims', f'{data_dim}', '100', '2', 
                 '--gradient_type', 'score',
                 '--num_variational_samples', '1',
                 '--epochs', '50',
                 '--learning_rate', '1e-4',
                 '--vr_coeff', '320',
                 '--cuda']

args_score_vr = parser_vr.parse_args(args_score_vr)
args_score_vr.cuda = args_score_vr.cuda and torch.cuda.is_available()
args_score_vr.device = torch.device("cuda" if args_score_vr.cuda else "cpu")

# Seed all random number generators for reproducibility of the runs
seed_everything(args_score_vr.seed)

# Initialise the model and the Adam (SGD) optimiser
model_score_vr = VAEwScoreVR(args_score_vr).to(args_score_vr.device)
optimizer_score_vr = optim.Adam(model_score_vr.parameters(), lr=args_score_vr.learning_rate)

# save_state(model_score_vr, optimizer_score_vr, 'VAE_score_vr.pt')

# Alternatively to training a model from scratch you can load a pre-trained model
# load_state(model_score_vr, optimizer_score_vr, 'VAE_score_vr.pt')

# Specify the hyperpameter choices
data_dim = train_data[0][0].numel()
args_score_vr_mc100 = ['--dec_layer_dims', '2', '100', f'{data_dim}',
                       '--enc_layer_dims', f'{data_dim}', '100', '2',
                       '--gradient_type', 'score',
                       '--num_variational_samples', '100',
                       '--epochs', '50',
                       '--learning_rate', '7e-4',
                       '--vr_coeff', '320',
                       '--cuda']

args_score_vr_mc100 = parser_vr.parse_args(args_score_vr_mc100)
args_score_vr_mc100.cuda = args_score_vr_mc100.cuda and torch.cuda.is_available()
args_score_vr_mc100.device = torch.device("cuda" if args_score_vr_mc100.cuda else "cpu")

# Seed all random number generators for reproducibility of the runs
seed_everything(args_score_vr_mc100.seed)

# Initialise the model and the Adam (SGD) optimiser
model_score_vr_mc100 = VAEwScoreVR(args_score_vr_mc100).to(args_score_vr_mc100.device)
optimizer_score_vr_mc100 = optim.Adam(model_score_vr_mc100.parameters(), lr=args_score_vr_mc100.learning_rate)

# save_state(model_score_vr_mc100, optimizer_score_vr_mc100, 'VAE_score_vr_mc100.pt')

# Alternatively to training a model from scratch you can load a pre-trained model
# load_state(model_score_vr_mc100, optimizer_score_vr_mc100, 'VAE_score_vr_mc100.pt')

# Specify the hyperpameter choices
data_dim = train_data[0][0].numel()
args_score_mc100 = ['--dec_layer_dims', '2', '100', f'{data_dim}',
                       '--enc_layer_dims', f'{data_dim}', '100', '2',
                       '--gradient_type', 'score',
                       '--num_variational_samples', '100',
                       '--epochs', '50',
                       '--learning_rate', '1e-4',
                       '--cuda']

args_score_mc100 = parser_vr.parse_args(args_score_mc100)
args_score_mc100.cuda = args_score_mc100.cuda and torch.cuda.is_available()
args_score_mc100.device = torch.device("cuda" if args_score_mc100.cuda else "cpu")

# Seed all random number generators for reproducibility of the runs
seed_everything(args_score_mc100.seed)

# Initialise the model and the Adam (SGD) optimiser
model_score_mc100 = VAE(args_score_mc100).to(args_score_mc100.device)
optimizer_score_mc100 = optim.Adam(model_score_mc100.parameters(), lr=args_score_mc100.learning_rate)

# save_state(model_score_mc100, optimizer_score_mc100, 'VAE_score_mc100.pt')

# Alternatively to training a model from scratch you can load a pre-trained model
# load_state(model_score_mc100, optimizer_score_mc100, 'VAE_score_mc100.pt')