<repo_name>henrhoi/made-pytorch<file_sep>/made.py
import torch.nn as nn
from linear_masked import LinearMasked
from utils import *


class MADE(nn.Module):
    """
    Class implementing MADE for 1 channel image-input based on "MADE: Masked Autoencoder for Distribution Estimation" by <NAME>. al.
    """

    def __init__(self, input_dim):
        super().__init__()
        self.input_dim = input_dim

        self.net = nn.Sequential(
            LinearMasked(input_dim, input_dim),
            nn.ReLU(),
            LinearMasked(input_dim, input_dim),
            nn.ReLU(),
            LinearMasked(input_dim, input_dim)
        )
        self.apply_masks()

    def forward(self, x):
        return self.net(x)

    def apply_masks(self):
        # Set order of masks, i.e. who can make which edges
        ordering = np.arange(self.input_dim)

        # Create masks
        masks = []
        masks.append(ordering[:, None] <= ordering[None, :])
        masks.append(ordering[:, None] <= ordering[None, :])
        masks.append(ordering[:, None] < ordering[None, :])

        # Set the masks in all LinearMasked layers
        layers = [layer for layer in self.net.modules() if isinstance(layer, LinearMasked)]
        for i in range(len(layers)):
            layers[i].set_mask(masks[i])


def train_made(train_data, test_data, image_shape, dset_id):
    """
    train_data: A (n_train, H, W, 1) uint8 numpy array of binary images with values in {0, 1}
    test_data: An (n_test, H, W, 1) uint8 numpy array of binary images with values in {0, 1}
    image_shape: (H, W), height and width of the image
    dset_id: An identifying number of which dataset is given (1 or 2). Most likely
            used to set different hyperparameters for different datasets

    Returns
    - a (# of training iterations,) numpy array of train_losses evaluated every minibatch
    - a (# of epochs + 1,) numpy array of test_losses evaluated once at initialization and after each epoch
    - a numpy array of size (100, H, W, 1) of samples with values in {0, 1}
    """

    use_cuda = False
    device = torch.device('cuda') if use_cuda else None

    train_data = torch.from_numpy(
        train_data.reshape((train_data.shape[0], train_data.shape[1] * train_data.shape[2]))).float().to(device)
    test_data = torch.from_numpy(
        test_data.reshape((test_data.shape[0], test_data.shape[1] * test_data.shape[2]))).float().to(device)

    def nll_loss(batch, output):
        return torch.nn.functional.binary_cross_entropy(torch.sigmoid(output), batch)

    H, W = image_shape[0], image_shape[1]
    input_dim = H * W

    dataset_params = {
        'batch_size': 32,
        'shuffle': True
    }

    made = MADE(input_dim).cuda() if use_cuda else MADE(input_dim)
    n_epochs = 10
    lr = 0.003

    train_loader = torch.utils.data.DataLoader(train_data, **dataset_params)
    optimizer = torch.optim.Adam(made.parameters(), lr=lr)

    init_test_loss = nll_loss(test_data, made(test_data))

    train_losses = []
    test_losses = [init_test_loss.item()]

    for epoch in range(n_epochs):
        for batch_x in train_loader:
            optimizer.zero_grad()
            output = made(batch_x)
            loss = nll_loss(batch_x, output)
            loss.backward()
            optimizer.step()
            train_losses.append(loss.item())

        test_loss = nll_loss(test_data, made(test_data))
        test_losses.append(test_loss.item())
        print(f"[{100*(epoch+1)/n_epochs:.2f}%] Epoch {epoch + 1}")

    if use_cuda:
        torch.cuda.empty_cache()

    samples = torch.zeros(size=(100, H * W)).to(device)
    made.eval()
    with torch.no_grad():
        for i in range(H * W):
            out = made(samples)
            proba = torch.sigmoid(out)
            torch.bernoulli(proba[:, i], out=samples[:, i])

    return np.array(train_losses), np.array(test_losses), samples.reshape((100, H, W, 1)).detach().cpu().numpy()


def train_and_show_results_shapes_data():
    """
    Trains MADE and displays samples and training plot for Shapes dataset
    """
    show_results_made(1, 'b', train_made)


def train_and_show_results_mnist_data():
    """
    Trains MADE and displays samples and training plot for MNIST dataset
    """
    show_results_made(2, 'b', train_made)
<file_sep>/made2d.py
import torch.nn as nn
from linear_masked import LinearMasked
from utils import *


class MADE2D(nn.Module):
    """
    Class implementing MADE for 2D input based on "MADE: Masked Autoencoder for Distribution Estimation" by <NAME>. al.
    """

    def __init__(self, input_dim, d):
        super().__init__()
        self.d = d
        self.input_dim = input_dim

        self.net = nn.Sequential(
            LinearMasked(input_dim, input_dim),
            nn.ReLU(),
            LinearMasked(input_dim, input_dim),
            nn.ReLU(),
            LinearMasked(input_dim, input_dim)
        )
        self.apply_masks()

    def forward(self, x):
        return self.net(x)

    def apply_masks(self):
        # Set order of masks, i.e. who can make which edges
        order_0 = np.concatenate((np.repeat(1, self.d), np.repeat(2, self.d)))
        order_1 = np.repeat(1, self.input_dim)
        order_2 = np.repeat(1, self.input_dim)

        # Create masks
        masks = []
        masks.append(order_0[:, None] <= order_1[None, :])
        masks.append(order_1[:, None] <= order_2[None, :])
        masks.append(order_2[:, None] < order_0[None, :])

        # Set the masks in all LinearMasked layers
        layers = [layer for layer in self.net.modules() if isinstance(layer, LinearMasked)]
        for i in range(len(layers)):
            layers[i].set_mask(masks[i])


def train_made(train_data, test_data, d, dset_id):
    """
    train_data: An (n_train, 2) numpy array of integers in {0, ..., d-1}
    test_data: An (n_test, 2) numpy array of integers in {0, .., d-1}
    d: The number of possible discrete values for each random variable x1 and x2
    dset_id: An identifying number of which dataset is given (1 or 2). Most likely
            used to set different hyperparameters for different datasets

    Returns
    - a (# of training iterations,) numpy array of train_losses evaluated every minibatch
    - a (# of epochs + 1,) numpy array of test_losses evaluated once at initialization and after each epoch
    - a numpy array of size (d, d) of probabilities (the learned joint distribution)
    """

    def make_one_hot(data, num_classes=d):
        one_hot = np.zeros((len(data), 2 * d))
        one_hot[np.arange(data.shape[0]), data[:, 0]] = 1
        one_hot[np.arange(data.shape[0]), d + data[:, 1]] = 1
        return torch.from_numpy(one_hot).float()

    use_cuda = True
    device = torch.device('cuda') if use_cuda else None

    train_data = make_one_hot(train_data).to(device)
    test_data = make_one_hot(test_data).to(device)

    def get_proba(batch, output, d):
        output_1 = torch.softmax(output[:, :d], dim=1)
        output_2 = torch.softmax(output[:, d:], dim=1)

        indices_1 = torch.nonzero(batch[:, :d])
        indices_2 = torch.nonzero(batch[:, d:])

        return torch.gather(input=output_1, dim=1, index=indices_1[:, [1]]) \
               * torch.gather(input=output_2, dim=1, index=indices_2[:, [1]])

    def nll_loss(batch, output, d):
        proba = get_proba(batch, output, d)
        return -torch.mean(torch.log(proba + 1e-8))

    made = MADE2D(2 * d, d).cuda() if use_cuda else MADE2D(2 * d, d)
    dataset_params = {
        'batch_size': 32,
        'shuffle': True
    }

    n_epochs = 50 if dset_id == 2 else 10
    lr = 0.0025 if dset_id == 2 else 0.01
    train_loader = torch.utils.data.DataLoader(train_data, **dataset_params)
    optimizer = torch.optim.Adam(made.parameters(), lr=lr)

    init_test_loss = nll_loss(test_data, made(test_data), d)

    train_losses = []
    test_losses = [init_test_loss.item()]

    for epoch in range(n_epochs):
        for batch_x in train_loader:
            optimizer.zero_grad()
            output = made(batch_x)
            loss = nll_loss(batch_x, output, d)
            loss.backward()
            optimizer.step()
            train_losses.append(loss.item())

        test_loss = nll_loss(torch.tensor(test_data), made(torch.tensor(test_data)), d)
        test_losses.append(test_loss.item())
        print(f"[{100*(epoch+1)/n_epochs:.2f}%] Epoch {epoch + 1}") if epoch % 2 else None

    x_s = make_one_hot(torch.tensor([[i, j] for i in range(d) for j in range(d)])).to(device)

    proba = get_proba(x_s, made(x_s), d).reshape(d, d)

    return np.array(train_losses), np.array(test_losses), proba.detach().cpu().numpy()


def train_and_show_results_smiley_data():
    """
    Trains 2D MADE and displays samples and training plot for smiley distribution
    """
    show_results_made(1, 'a', train_made)


def train_and_show_results_geoffrey_hinton_data():
    """
    Trains 2D MADE and displays samples and training plot for Geoffrey Hinton distribution
    """
    show_results_made(2, 'a', train_made)
<file_sep>/utils.py
import pickle
from os.path import join
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
import numpy as np

# Mixture of logistics
import torch


def mixture_of_logistics_sample_data_1():
    count = 1000
    rand = np.random.RandomState(0)
    samples = 0.4 + 0.1 * rand.randn(count)
    data = np.digitize(samples, np.linspace(0.0, 1.0, 20))
    split = int(0.8 * len(data))
    train_data, test_data = data[:split], data[split:]
    return train_data, test_data


def mixture_of_logistics_sample_data_2():
    count = 10000
    rand = np.random.RandomState(0)
    a = 0.3 + 0.1 * rand.randn(count)
    b = 0.8 + 0.05 * rand.randn(count)
    mask = rand.rand(count) < 0.5
    samples = np.clip(a * mask + b * (1 - mask), 0.0, 1.0)
    data = np.digitize(samples, np.linspace(0.0, 1.0, 100))
    split = int(0.8 * len(data))
    train_data, test_data = data[:split], data[split:]
    return train_data, test_data


def show_training_plot(train_losses, test_losses, title):
    plt.figure()
    n_epochs = len(test_losses) - 1
    x_train = np.linspace(0, n_epochs, len(train_losses))
    x_test = np.arange(n_epochs + 1)

    plt.plot(x_train, train_losses, label='train loss')
    plt.plot(x_test, test_losses, label='test loss')
    plt.legend()
    plt.title(title)
    plt.xlabel('Epoch')
    plt.ylabel('NLL')
    plt.show()


def show_distribution_1d(data, distribution, title):
    d = len(distribution)

    plt.figure()
    plt.hist(data, bins=np.arange(d) - 0.5, label='train data', density=True)

    x = np.linspace(-0.5, d - 0.5, 1000)
    y = distribution.repeat(1000 // d)
    plt.plot(x, y, label='learned distribution')

    plt.title(title)
    plt.xlabel('x')
    plt.ylabel('Probability')
    plt.legend()
    plt.show()


def show_results(dset_type, fn):
    if dset_type == 1:
        train_data, test_data = mixture_of_logistics_sample_data_1()
        d = 20
    elif dset_type == 2:
        train_data, test_data = mixture_of_logistics_sample_data_2()
        d = 100
    else:
        raise Exception('Invalid dset_type:', dset_type)

    train_losses, test_losses, distribution = fn(train_data, test_data, d, dset_type)
    assert np.allclose(np.sum(distribution), 1), f'Distribution sums to {np.sum(distribution)} != 1'

    print(f'Final Test Loss: {test_losses[-1]:.4f}')

    show_training_plot(train_losses, test_losses, f'Dataset {dset_type} Train Plot')
    show_distribution_1d(train_data, distribution, f'Dataset {dset_type} Learned Distribution')


# MADE
def load_pickled_data(fname, include_labels=False):
    with open(fname, 'rb') as f:
        data = pickle.load(f)

    train_data, test_data = data['train'], data['test']
    if 'mnist.pkl' in fname or 'shapes.pkl' in fname:
        # Binarize MNIST and shapes dataset
        train_data = (train_data > 127.5).astype('uint8')
        test_data = (test_data > 127.5).astype('uint8')
    if 'celeb.pkl' in fname:
        train_data = train_data[:, :, :, [2, 1, 0]]
        test_data = test_data[:, :, :, [2, 1, 0]]
    if include_labels:
        return train_data, test_data, data['train_labels'], data['test_labels']
    return train_data, test_data


# Question 2
def sample_2d_data(image_file, n, d):
    from PIL import Image
    import itertools

    im = Image.open(image_file).resize((d, d)).convert('L')
    im = np.array(im).astype('float32')
    dist = im / im.sum()

    pairs = list(itertools.product(range(d), range(d)))
    idxs = np.random.choice(len(pairs), size=n, replace=True, p=dist.reshape(-1))
    samples = [pairs[i] for i in idxs]

    return dist, np.array(samples)


def show_distribution_2d(true_dist, learned_dist):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8))
    ax1.imshow(true_dist)
    ax1.set_title('True Distribution')
    ax1.axis('off')
    ax2.imshow(learned_dist)
    ax2.set_title('Learned Distribution')
    ax2.axis('off')
    plt.show()


def show_samples(samples, nrow=10, title='Samples'):
    samples = (torch.FloatTensor(samples) / 255).permute(0, 3, 1, 2)
    grid_img = make_grid(samples, nrow=nrow)
    plt.figure()
    plt.title(title)
    plt.imshow(grid_img.permute(1, 2, 0))
    plt.axis('off')
    plt.show()


def show_results_made(dset_type, part, fn):
    data_dir = "data"
    if part == 'a':
        if dset_type == 1:
            n, d = 10000, 25
            true_dist, data = sample_2d_data(join(data_dir, 'smiley.jpg'), n, d)
        elif dset_type == 2:
            n, d = 100000, 200
            true_dist, data = sample_2d_data(join(data_dir, 'geoffrey-hinton.jpg'), n, d)
        else:
            raise Exception('Invalid dset_type:', dset_type)
        split = int(0.8 * len(data))
        train_data, test_data = data[:split], data[split:]
    elif part == 'b':
        if dset_type == 1:
            train_data, test_data = load_pickled_data(join(data_dir, 'shapes.pkl'))
            img_shape = (20, 20)
        elif dset_type == 2:
            train_data, test_data = load_pickled_data(join(data_dir, 'mnist.pkl'))
            img_shape = (28, 28)
        else:
            raise Exception('Invalid dset type:', dset_type)
    else:
        raise Exception('Invalid part', part)

    if part == 'a':
        train_losses, test_losses, distribution = fn(train_data, test_data, d, dset_type)
        assert np.allclose(np.sum(distribution), 1), f'Distribution sums to {np.sum(distribution)} != 1'

        print(f'Final Test Loss: {test_losses[-1]:.4f}')

        show_training_plot(train_losses, test_losses, f'Dataset {dset_type} Train Plot)')
        show_distribution_2d(true_dist, distribution)

    elif part == 'b':
        train_losses, test_losses, samples = fn(train_data, test_data, img_shape, dset_type)
        samples = samples.astype('float32') * 255
        print(f'Final Test Loss: {test_losses[-1]:.4f}')
        show_training_plot(train_losses, test_losses, f'Dataset {dset_type} Train Plot')
        show_samples(samples)
<file_sep>/mixture_of_logistics.py
from utils import *


def generative_model_histogram(train_data, test_data, d, dset_id):
    """
    train_data: An (n_train,) numpy array of integers in {0, ..., d-1}
    test_data: An (n_test,) numpy array of integers in {0, .., d-1}
    d: The number of possible discrete values for random variable x
    dset_id: An identifying number of which dataset is given (1 or 2). Most likely
                used to set different hyperparameters for different datasets

    Returns
    - a (# of training iterations,) numpy array of train_losses evaluated every minibatch
    - a (# of epochs + 1,) numpy array of test_losses evaluated once at initialization and after each epoch
    - a numpy array of size (d,) of model probabilities
    """

    def nll_loss(x, theta):
        model = torch.exp(torch.gather(theta, dim=0, index=x)) / torch.exp(theta).sum()
        return torch.mean(-torch.log(model))

    dataset_params = {
        'batch_size': 32,
        'shuffle': True
    }

    theta = torch.zeros(d, dtype=torch.float32, requires_grad=True)
    n_epochs = 40
    train_loader = torch.utils.data.DataLoader(train_data, **dataset_params)
    optimizer = torch.optim.Adam([theta], lr=0.01)

    train_losses = []
    test_losses = [nll_loss(torch.tensor(test_data), theta).item()]

    for epoch in range(n_epochs):
        for batch_x in train_loader:
            optimizer.zero_grad()
            loss = nll_loss(batch_x, theta)
            loss.backward()
            optimizer.step()
            train_losses.append(loss.item())

        test_loss = nll_loss(torch.tensor(test_data), theta)
        test_losses.append(test_loss.item())

    return np.array(train_losses), np.array(test_losses), torch.softmax(theta, dim=0).detach().numpy()


def discretized_mixture_of_logistics(train_data, test_data, d, dset_id):
    """
    train_data: An (n_train,) numpy array of integers in {0, ..., d-1}
    test_data: An (n_test,) numpy array of integers in {0, .., d-1}
    d: The number of possible discrete values for random variable x
    dset_id: An identifying number of which dataset is given (1 or 2). Most likely
            used to set different hyperparameters for different datasets

    Returns
    - a (# of training iterations,) numpy array of train_losses evaluated every minibatch
    - a (# of epochs + 1,) numpy array of test_losses evaluated once at initialization and after each epoch
    - a numpy array of size (d,) of model probabilities
    """

    def model(x, pis, means, stds):
        x = x.reshape((x.shape[0], 1))
        x_plus = x + 0.5
        x_minus = x - 0.5

        # Handeling edge cases
        x_plus[x_plus > d - 1] = torch.tensor(float('inf'), requires_grad=True)
        x_minus[x_minus < 0] = torch.tensor(float('-inf'), requires_grad=True)

        unweighted_mixtures = torch.sigmoid((x_plus - means) / torch.exp(stds)) - torch.sigmoid(
            (x_minus - means) / torch.exp(stds))
        model_proba = (unweighted_mixtures * torch.softmax(pis, dim=1)).sum(dim=1, keepdim=False)

        return model_proba

    def nll_loss(x, pis, means, stds):
        model_proba = model(x, pis, means, stds)
        return torch.mean(-torch.log(model_proba))

    dataset_params = {
        'batch_size': 32,
        'shuffle': True
    }

    if dset_id == 2:
        pis = torch.tensor([[1., 2., 3., 4.]], requires_grad=True)
        means = torch.tensor([[20., 30., 80., 90.]], requires_grad=True)
        stds = torch.tensor([[1., 2., 3., 4.]], requires_grad=True)
    else:
        pis = torch.tensor([[1., 2., 3., 4.]], requires_grad=True)
        means = torch.tensor([[1., 2., 3., 4.]], requires_grad=True)
        stds = torch.tensor([[1., 2., 3., 4.]], requires_grad=True)

    n_epochs = 20
    train_loader = torch.utils.data.DataLoader(train_data, **dataset_params)
    optimizer = torch.optim.Adam([pis, means, stds], lr=0.01)

    train_losses = []
    test_losses = [nll_loss(torch.tensor(test_data), pis, means, stds).item()]

    for epoch in range(n_epochs):
        for batch_x in train_loader:
            optimizer.zero_grad()
            loss = nll_loss(batch_x, pis, means, stds)
            loss.backward()
            optimizer.step()
            train_losses.append(loss.item())

        test_loss = nll_loss(torch.tensor(test_data), pis, means, stds)
        test_losses.append(test_loss.item())

    x_s = torch.tensor([i for i in range(d)])
    model_probabilities = model(x_s, pis, means, stds)
    return np.array(train_losses), np.array(test_losses), model_probabilities.detach().numpy()


if __name__ == '__main__':
    show_results(dset_type=1, fn=generative_model_histogram)
    show_results(dset_type=2, fn=discretized_mixture_of_logistics)


<file_sep>/linear_masked.py
import torch
import torch.nn as nn


class LinearMasked(nn.Linear):
    """
    Class implementing nn.Linear with mask
    """

    def __init__(self, in_features, out_features, bias=True):
        super().__init__(in_features, out_features, bias)
        self.register_buffer('mask', torch.ones(out_features, in_features))

    def set_mask(self, mask):
        self.mask.data.copy_(torch.from_numpy(mask.astype(np.uint8).T))

    def forward(self, x):
        return torch.nn.functional.linear(x, self.mask * self.weight, self.bias)
<file_sep>/README.md
# MADE (Masked Autoencoder Distribution Estimation) in PyTorch
PyTorch implementations of MADE (Masked Autoencoder Distribution Estimation) based on [MADE: Masked Autoencoder for Distribution Estimation](https://arxiv.org/abs/1502.03509) by Germain et. al. and autoregressive models.

## Models

**MADE - Fitting 2D Data:**

Implemented a [MADE](https://arxiv.org/abs/1502.03509) model through maximum likelihood to represent  <img src="https://render.githubusercontent.com/render/math?math=p(x_0,x_1)">  on the given datasets, with any autoregressive ordering.

**MADE:**

Implemented a MADE model on the binary image datasets Shape and MNIST. Given some binary image of height <img src="https://render.githubusercontent.com/render/math?math=H"> and width <img src="https://render.githubusercontent.com/render/math?math=W">, we can represent image <img src="https://render.githubusercontent.com/render/math?math=x\in \{0, 1\}^{H\times W}"> as a flattened binary vector <img src="https://render.githubusercontent.com/render/math?math=x\in \{0, 1\}^{HW}"> to input into MADE to model <img src="https://render.githubusercontent.com/render/math?math=p_\theta(x) = \prod_{i=1}^{HW} p_\theta(x_i|x_{ < i})">.

Illustration of MADE-model:

![made](https://i.imgur.com/agJN65k.png)

**Fitting a histogram:**

Let <img src="https://render.githubusercontent.com/render/math?math=\theta = (\theta_0, \dots, \theta_{d-1}) \in \mathbb{R}^d"> and define the model <img src="https://render.githubusercontent.com/render/math?math=p_\theta(x) = \dfrac{e^{\theta_x}}{\sum_{x'}e^{\theta_{x'}}}">

Fit <img src="https://render.githubusercontent.com/render/math?math=p_\theta"> with maximum likelihood via stochastic gradient descent on the training set, using <img src="https://render.githubusercontent.com/render/math?math=\theta"> initialized to zero.

**Discretized Mixture of Logistics:**

Let us model <img src="https://render.githubusercontent.com/render/math?math=p_\theta(x)"> as a **discretized** mixture of 4 logistics such that <img src="https://render.githubusercontent.com/render/math?math=p_\theta(x) = \sum_{i=1}^4 \pi_i[\sigma((x+0.5 - \mu_i)/s_i) - \sigma((x-0.5-\mu_i)/s_i)]">


For the edge case of when <img src="https://render.githubusercontent.com/render/math?math=x = 0">, we replace <img src="https://render.githubusercontent.com/render/math?math=x-0.5"> by <img src="https://render.githubusercontent.com/render/math?math=-\infty">, and for <img src="https://render.githubusercontent.com/render/math?math=x = d-1">, we replace <img src="https://render.githubusercontent.com/render/math?math=x+0.5"> by <img src="https://render.githubusercontent.com/render/math?math=\infty">.

One may find the [PixelCNN++](https://arxiv.org/abs/1701.05517) helpful for more information on discretized mixture of logistics.

## Results and samples

| Model | Dataset |  Result/Samples |
|:---:|:---:|:---:|
| MADE 2D (Smiley)                | ![3](https://i.imgur.com/OLtbWDr.png)  | ![33](https://i.imgur.com/En6OGv4.png)  |
| MADE 2D (<NAME>)       | ![4](https://i.imgur.com/c8PYEtz.png)  | ![44](https://i.imgur.com/kvPCx4y.png)  |
| MADE (Shapes)                   | ![5](https://i.imgur.com/QGWaShQ.png)  | ![55](https://i.imgur.com/o0z8lEC.png)  |
| MADE (MNIST)                    | ![6](https://i.imgur.com/T0acOeg.png)  | ![66](https://i.imgur.com/pMjtBbj.png)  |
| Simple histogram                | ![1](https://i.imgur.com/jPCPDmD.png)  | ![11](https://i.imgur.com/MLujTCa.png)  |
| Discretized Mixture of Logistics| ![2](https://i.imgur.com/jPCPDmD.png)  | ![22](https://i.imgur.com/OthJ9Gf.png)  |
