add app and generation / model code

app.py

@@ -0,0 +1,34 @@
+import torch
+import torchvision
+import gradio as gr
+from PIL import Image
+from cli import iterative_refinement
+from viz import grid_of_images_default
+from subprocess
+subprocess.call("download_models.sh", shell=True)
+models = {
+    "convae": torch.load("convae.th", map_location="cpu"),
+    "deep_convae": torch.load("deep_convae.th", map_location="cpu"),
+}
+
+def gen(model, seed, nb_iter, nb_samples, width, height):
+    torch.manual_seed(int(seed))
+    bs = 64
+    model = models[model]
+    samples = iterative_refinement(
+        model, 
+        nb_iter=int(nb_iter),
+        nb_examples=int(nb_samples), 
+        w=int(width), h=int(height), c=1, 
+        batch_size=bs,
+    )
+    grid = grid_of_images_default(samples.reshape((samples.shape[0]*samples.shape[1], int(height), int(width), 1)).numpy(), shape=(samples.shape[0], samples.shape[1])) 
+    grid = (grid*255).astype("uint8")
+    return Image.fromarray(grid)
+
+iface = gr.Interface(
+    fn=gen,
+    inputs=[gr.Dropdown(list(models.keys()), value="deep_convae"), gr.Number(value=0), gr.Number(value=20), gr.Number(value=1), gr.Number(value=28), gr.Number(value=28)],
+    outputs="image"
+)
+iface.launch()

cli.py

@@ -0,0 +1,320 @@
+import os
+import matplotlib as mpl
+mpl.use('Agg')
+import matplotlib.pyplot as plt
+from functools import partial
+
+from clize import run
+import numpy as np
+from skimage.io import imsave
+
+from viz import grid_of_images_default
+
+import torch.nn as nn
+import torch
+
+from model import DenseAE
+from model import ConvAE
+from model import DeepConvAE
+from model import SimpleConvAE
+from model import ZAE
+from model import KAE
+from data import load_dataset
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def plot_dataset(code_2d, categories):
+    colors = [
+        'r',
+        'b',
+        'g',
+        'crimson',
+        'gold',
+        'yellow',
+        'maroon',
+        'm',
+        'c',
+        'orange'
+    ]
+    for cat in range(0, 10):
+        g = (categories == cat)
+        plt.scatter(
+            code_2d[g, 0], 
+            code_2d[g, 1],
+            marker='+', 
+            c=colors[cat], 
+            s=40, 
+            alpha=0.7,
+            label="digit {}".format(cat)
+        )
+
+
+def plot_generated(code_2d, categories):
+    g = (categories < 0)
+    plt.scatter(
+        code_2d[g, 0], 
+        code_2d[g, 1], 
+        marker='+',
+        c='gray', 
+        s=30
+    )
+
+
+def grid_embedding(h):
+    from lapjv import lapjv
+    from scipy.spatial.distance import cdist
+    assert int(np.sqrt(h.shape[0])) ** 2 == h.shape[0], 'Nb of examples must be a square number'
+    size = int(np.sqrt(h.shape[0]))
+    grid = np.dstack(np.meshgrid(np.linspace(0, 1, size), np.linspace(0, 1, size))).reshape(-1, 2)
+    cost_matrix = cdist(grid, h, "sqeuclidean").astype('float32')
+    cost_matrix = cost_matrix * (100000 / cost_matrix.max())
+    _, rows, cols = lapjv(cost_matrix)
+    return rows
+
+
+def save_weights(m, folder='.'):
+    if isinstance(m, nn.Linear):
+        w = m.weight.data
+        if w.size(1) == 28*28 or w.size(0) == 28*28:
+            w0, w1 = w.size(0), w.size(1)
+            if w0 == 28*28:
+                w = w.transpose(0, 1)
+                w = w.contiguous()
+            w = w.view(w.size(0), 1, 28, 28)
+            gr = grid_of_images_default(np.array(w.tolist()), normalize=True)
+            imsave('{}/feat_{}.png'.format(folder, w0), gr)
+    elif isinstance(m, nn.ConvTranspose2d):
+        w = m.weight.data
+        if w.size(0) in (32, 64, 128, 256, 512) and w.size(1) in (1, 3):
+            gr = grid_of_images_default(np.array(w.tolist()), normalize=True)
+            imsave('{}/feat.png'.format(folder), gr)
+
+@torch.no_grad()
+def iterative_refinement(ae, nb_examples=1, nb_iter=10, w=28, h=28, c=1, batch_size=None):
+    if batch_size is None:
+        batch_size = nb_examples
+    x = torch.rand(nb_iter, nb_examples, c, w, h)
+    for i in range(1, nb_iter):
+        for j in range(0, nb_examples, batch_size):
+            oldv = x[i-1][j:j + batch_size].to(device)
+            newv = ae(oldv)
+            newv = newv.data.cpu()
+            x[i][j:j + batch_size] = newv
+    return x
+
+
+def build_model(name, w, h, c):
+    if name == 'convae':
+        ae = ConvAE(
+            w=w, h=h, c=c, 
+            nb_filters=128, 
+            spatial=True, 
+            channel=True, 
+            channel_stride=4,
+        )
+    elif name == 'zae':
+        ae = ZAE(
+            w=w, h=h, c=c,
+            theta=3,
+            nb_hidden=1000,
+        )
+    elif name == 'kae':
+        ae = KAE(
+            w=w, h=h, c=c,
+            nb_active=1000,
+            nb_hidden=1000,
+        )
+    elif name == 'denseae':
+        ae = DenseAE(
+            w=w, h=h, c=c,
+            encode_hidden=[1000],
+            decode_hidden=[],
+            ksparse=True,
+            nb_active=50,
+        )
+    elif name == 'simple_convae':
+        ae = SimpleConvAE(
+            w=w, h=h, c=c,
+            nb_filters=128,
+        )
+    elif name == 'deep_convae':
+        ae = DeepConvAE(
+            w=w, h=h, c=c, 
+            nb_filters=128, 
+            spatial=True, 
+            channel=True, 
+            channel_stride=4,
+            nb_layers=3, 
+        )
+    else:
+        raise ValueError('Unknown model')
+
+    return ae
+
+
+def salt_and_pepper(X, proba=0.5):
+    a = (torch.rand(X.size()).to(device) <= (1 - proba)).float()
+    b = (torch.rand(X.size()).to(device) <= 0.5).float()
+    c = ((a == 0).float() * b)
+    return X * a + c
+
+
+def train(*, dataset='mnist', folder='mnist', resume=False, model='convae', walkback=False, denoise=False, epochs=100, batch_size=64, log_interval=100):
+    gamma = 0.99
+    dataset = load_dataset(dataset, split='train')
+    x0, _ = dataset[0]
+    c, h, w = x0.size()
+    dataloader = torch.utils.data.DataLoader(
+        dataset, 
+        batch_size=batch_size,
+        shuffle=True, 
+        num_workers=4
+    )
+    if resume:
+        ae = torch.load('{}/model.th'.format(folder))
+        ae = ae.to(device)
+    else:
+        ae = build_model(model, w=w, h=h, c=c)
+        ae = ae.to(device)
+    optim = torch.optim.Adadelta(ae.parameters(), lr=0.1, eps=1e-7, rho=0.95, weight_decay=0)
+    avg_loss = 0.
+    nb_updates = 0
+    _save_weights = partial(save_weights, folder=folder)
+
+    for epoch in range(epochs):
+        for X, y in dataloader:
+            ae.zero_grad()
+            X = X.to(device)
+            if hasattr(ae, 'nb_active'):
+                ae.nb_active = max(ae.nb_active - 1, 32)
+            # walkback + denoise
+            if walkback:
+                loss = 0.
+                x = X.data
+                nb = 5
+                for _ in range(nb):
+                    x = salt_and_pepper(x, proba=0.3) # denoise
+                    x = x.to(device)
+                    x = ae(x) # reconstruct
+                    Xr = x
+                    loss += (((x - X) ** 2).view(X.size(0), -1).sum(1).mean()) / nb
+                    x = (torch.rand(x.size()).to(device) <= x.data).float() # sample
+            # denoise only
+            elif denoise:
+                Xc = salt_and_pepper(X.data, proba=0.3)
+                Xr = ae(Xc)
+                loss = ((Xr - X) ** 2).view(X.size(0), -1).sum(1).mean()
+            # normal training
+            else:
+                Xr = ae(X)
+                loss = ((Xr - X) ** 2).view(X.size(0), -1).sum(1).mean()
+            loss.backward()
+            optim.step()
+            avg_loss = avg_loss * gamma + loss.item() * (1 - gamma)
+            if nb_updates % log_interval == 0:
+                print('Epoch : {:05d} AvgTrainLoss: {:.6f}, Batch Loss : {:.6f}'.format(epoch, avg_loss, loss.item()  ))
+                gr = grid_of_images_default(np.array(Xr.data.tolist()))
+                imsave('{}/rec.png'.format(folder), gr)
+                ae.apply(_save_weights)
+                torch.save(ae, '{}/model.th'.format(folder))
+            nb_updates += 1
+
+
+def test(*, dataset='mnist', folder='out', model_path=None, nb_iter=100, nb_generate=100, tsne=False):
+    if not os.path.exists(folder):
+        os.makedirs(folder, exist_ok=True)
+    dataset = load_dataset(dataset, split='train')
+    x0, _ = dataset[0]
+    c, h, w = x0.size()
+    nb = nb_generate
+    print('Load model...')
+    if model_path is None:
+        model_path = os.path.join(folder, "model.th")
+    ae = torch.load(model_path, map_location="cpu")
+    ae = ae.to(device)
+    def enc(X):
+        batch_size = 64
+        h_list = []
+        for i in range(0, X.size(0), batch_size):
+            x = X[i:i + batch_size]
+            x = x.to(device)
+            name = ae.__class__.__name__
+            if name in ('ConvAE',):
+                h = ae.encode(x)
+                h, _ = h.max(2)
+                h = h.view((h.size(0), -1))
+            elif name in ('DenseAE',):
+                x = x.view(x.size(0), -1)
+                h = x
+                #h = ae.encode(x)
+            else:
+                h = x.view(x.size(0), -1)
+            h = h.data.cpu()
+            h_list.append(h)
+        return torch.cat(h_list, 0)
+
+    print('iterative refinement...')
+    g = iterative_refinement(
+        ae, 
+        nb_iter=nb_iter, 
+        nb_examples=nb, 
+        w=w, h=h, c=c, 
+        batch_size=64
+    )
+    np.savez('{}/generated.npz'.format(folder), X=g.numpy())
+    g_subset = g[:, 0:100]
+    gr = grid_of_images_default(g_subset.reshape((g_subset.shape[0]*g_subset.shape[1], h, w, 1)).numpy(), shape=(g_subset.shape[0], g_subset.shape[1])) 
+    imsave('{}/gen_full_iters.png'.format(folder), gr)
+
+    g = g[-1] # last iter
+    print(g.shape)
+    gr = grid_of_images_default(g.numpy())
+    imsave('{}/gen_full.png'.format(folder), gr)
+
+    if tsne:
+        from sklearn.manifold import TSNE
+        dataloader = torch.utils.data.DataLoader(
+            dataset, 
+            batch_size=nb,
+            shuffle=True, 
+            num_workers=1
+        )
+        print('Load data...')
+        X, y = next(iter(dataloader))
+        print('Encode data...')
+        xh = enc(X)
+        print('Encode generated...')
+        gh = enc(g)
+        X = X.numpy()
+        g = g.numpy()
+        xh = xh.numpy()
+        gh = gh.numpy()
+
+        a = np.concatenate((X, g), axis=0)
+        ah = np.concatenate((xh, gh), axis=0)
+        labels = np.array(y.tolist() + [-1] * len(g))
+        sne = TSNE()
+        print('fit tsne...')
+        ah = sne.fit_transform(ah)
+        print('grid embedding...')
+        
+        asmall = np.concatenate((a[0:450], a[nb:nb + 450]), axis=0)
+        ahsmall = np.concatenate((ah[0:450], ah[nb:nb + 450]), axis=0)
+        rows = grid_embedding(ahsmall)
+        asmall = asmall[rows]
+        gr = grid_of_images_default(asmall)
+        imsave('{}/sne_grid.png'.format(folder), gr)
+
+        fig = plt.figure(figsize=(10, 10))
+        plot_dataset(ah, labels)
+        plot_generated(ah, labels)
+        plt.legend(loc='best')
+        plt.axis('off')
+        plt.savefig('{}/sne.png'.format(folder))
+        plt.close(fig)
+
+
+
+if __name__ == '__main__':
+    run([train, test])

convert.py

@@ -0,0 +1,52 @@
+import numpy as np
+import torch, h5py
+from model import *
+w, h, c = 28, 28, 1
+model_new = DeepConvAE(
+    w=w, h=h, c=c, 
+    nb_filters=128, 
+    spatial=True, 
+    channel=True, 
+    channel_stride=4,
+    # total layers = nb_layers*2, where we have nb_layers for encoder and nb_layers for decoder
+    nb_layers=3, 
+)
+# model_old = h5py.File("mnist_deepconvae/model.h5")
+model_old = h5py.File("/home/mehdi/work/code/out_of_class/ae/mnist/model.h5")
+
+
+print(model_new)
+print(model_old["model_weights"].keys())
+
+
+for name, param in model_new.named_parameters():
+    enc_or_decode, layer_id, bias_or_kernel = name.split(".")
+
+    if enc_or_decode == "encode":
+        layer_name = "conv2d"
+    else:
+        layer_name = "up_conv2d"
+
+    layer_id = (int(layer_id)//2) + 1
+
+    full_layer_name = f"{layer_name}_{layer_id}"
+    print(full_layer_name)
+
+    k = "kernel" if bias_or_kernel == "weight" else "bias"
+    weights =  model_old["model_weights"][full_layer_name][full_layer_name][k][()]
+    weights = np.array(weights)
+    weights = torch.from_numpy(weights)
+    print(name, layer_id, param.shape, weights.shape)
+    inds = [4,3,2,1,0]
+    if k == "kernel":
+        if layer_name == "conv2d":
+            weights = weights.permute((3,2,0,1))
+            weights = weights[:,:,inds]
+            weights = weights[:,:,:, inds]
+            print("W", weights.shape)
+        elif layer_name == "up_conv2d":
+            weights = weights.permute((2,3,0,1))
+    print(param.shape, weights.shape)
+    param.data.copy_(weights)
+    print((param-weights).sum())
+torch.save(model_new, "mnist_deepconvae/model.th")

data.py

@@ -0,0 +1,94 @@
+import torch
+
+import torchvision.transforms as transforms
+import torchvision.datasets as dset
+
+
+class Invert:
+    def __call__(self, x):
+        return 1 - x
+
+class Gray:
+    def __call__(self, x):
+        return x[0:1]
+
+
+
+def load_dataset(dataset_name, split='full'):
+    if dataset_name == 'mnist':
+        dataset = dset.MNIST(
+            root='data/mnist', 
+            download=True,
+            transform=transforms.Compose([
+                transforms.ToTensor(),
+            ])
+        )
+        return dataset
+    elif dataset_name == 'coco':
+        dataset = dset.ImageFolder(root='data/coco',
+            transform=transforms.Compose([
+            transforms.Scale(64),
+            transforms.CenterCrop(64),
+            transforms.ToTensor(),
+         ]))
+        return dataset
+    elif dataset_name == 'quickdraw':
+        X = (np.load('data/quickdraw/teapot.npy'))
+        X = X.reshape((X.shape[0], 28, 28))
+        X  = X / 255.
+        X = X.astype(np.float32)
+        X = torch.from_numpy(X)
+        dataset = TensorDataset(X, X)
+        return dataset
+    elif dataset_name == 'shoes':
+        dataset = dset.ImageFolder(root='data/shoes/ut-zap50k-images/Shoes',
+            transform=transforms.Compose([
+            transforms.Scale(64),
+            transforms.CenterCrop(64),
+            transforms.ToTensor(),
+         ]))
+        return dataset
+    elif dataset_name == 'footwear':
+        dataset = dset.ImageFolder(root='data/shoes/ut-zap50k-images',
+            transform=transforms.Compose([
+            transforms.Scale(64),
+            transforms.CenterCrop(64),
+            transforms.ToTensor(),
+         ]))
+        return dataset
+    elif dataset_name == 'celeba':
+        dataset = dset.ImageFolder(root='data/celeba',
+            transform=transforms.Compose([
+            transforms.Scale(32),
+            transforms.CenterCrop(32),
+            transforms.ToTensor(),
+         ]))
+        return dataset
+    elif dataset_name == 'birds':
+        dataset = dset.ImageFolder(root='data/birds/'+split,
+            transform=transforms.Compose([
+            transforms.Scale(32),
+            transforms.CenterCrop(32),
+            transforms.ToTensor(),
+         ]))
+        return dataset
+    elif dataset_name == 'sketchy':
+        dataset = dset.ImageFolder(root='data/sketchy/'+split,
+            transform=transforms.Compose([
+            transforms.Scale(64),
+            transforms.CenterCrop(64),
+            transforms.ToTensor(),
+            Gray()
+         ]))
+        return dataset
+ 
+    elif dataset_name == 'fonts':
+        dataset = dset.ImageFolder(root='data/fonts/'+split,
+            transform=transforms.Compose([
+            transforms.ToTensor(),
+            Invert(),
+            Gray(),
+         ]))
+        return dataset
+    else:
+        raise ValueError('Error : unknown dataset')

model.py

@@ -0,0 +1,260 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.init import xavier_uniform
+
+class KAE(nn.Module):
+
+    def __init__(self, w=32, h=32, c=1, nb_hidden=300, nb_active=16):
+        super().__init__()
+        self.nb_hidden = nb_hidden
+        self.nb_active = nb_active
+        self.encode = nn.Sequential(
+            nn.Linear(w*h*c, nb_hidden, bias=False)
+        )
+        self.bias = nn.Parameter(torch.zeros(w*h*c))
+        self.params = nn.ParameterList([self.bias])
+        self.apply(_weights_init)
+
+    def forward(self, X):
+        size = X.size()
+        X = X.view(X.size(0), -1)
+        h = self.encode(X)
+        Xr, _ = self.decode(h)
+        Xr = Xr.view(size)
+        return Xr
+    
+    def decode(self, h):
+        thetas, _ = torch.sort(h, dim=1, descending=True)
+        thetas = thetas[:, self.nb_active:self.nb_active+1]
+        h = h * (h > thetas).float()
+        Xr = torch.matmul(h, self.encode[0].weight) + self.bias
+        Xr = nn.Sigmoid()(Xr)
+        return Xr, h
+
+
+class ZAE(nn.Module):
+
+    def __init__(self, w=32, h=32, c=1, nb_hidden=300, theta=1):
+        super().__init__()
+        self.nb_hidden = nb_hidden
+        self.theta = theta
+        self.encode = nn.Sequential(
+            nn.Linear(w*h*c, nb_hidden, bias=False)
+        )
+        self.bias = nn.Parameter(torch.zeros(w*h*c))
+        self.params = nn.ParameterList([self.bias])
+        self.apply(_weights_init)
+
+    def forward(self, X):
+        size = X.size()
+        X = X.view(X.size(0), -1)
+        h = self.encode(X)
+        Xr, _ = self.decode(h)
+        Xr = Xr.view(size)
+        return Xr
+    
+    def decode(self, h):
+        h  = h * (h > self.theta).float()
+        Xr = torch.matmul(h, self.encode[0].weight) + self.bias
+        Xr = nn.Sigmoid()(Xr)
+        return Xr, h
+
+
+
+class DenseAE(nn.Module):
+
+    def __init__(self, w=32, h=32, c=1, encode_hidden=(300,), decode_hidden=(300,), ksparse=True, nb_active=10, denoise=None):
+        super().__init__()
+        self.encode_hidden = encode_hidden
+        self.decode_hidden = decode_hidden
+        self.ksparse = ksparse
+        self.nb_active = nb_active
+        self.denoise = denoise
+        
+        # encode layers
+        layers = []
+        hid_prev = w * h * c
+        for hid in encode_hidden:
+            layers.extend([
+                nn.Linear(hid_prev, hid),
+                nn.ReLU(True)
+            ])
+            hid_prev = hid
+        self.encode = nn.Sequential(*layers)
+        
+        # decode layers
+        layers = []
+        for hid in decode_hidden:
+            layers.extend([
+                nn.Linear(hid_prev, hid),
+                nn.ReLU(True)
+            ])
+            hid_prev = hid
+        layers.extend([
+            nn.Linear(hid_prev, w * h * c),
+            nn.Sigmoid()
+        ])
+        self.decode = nn.Sequential(*layers)
+        
+        self.apply(_weights_init)
+
+    def forward(self, X):
+        size = X.size()
+        if self.denoise is not None:
+            X = X * ((torch.rand(X.size()) <= self.denoise).float()).to(X.device)
+        X = X.view(X.size(0), -1)
+        h = self.encode(X)
+        if self.ksparse:
+            h = ksparse(h, nb_active=self.nb_active)
+        Xr = self.decode(h)
+        Xr = Xr.view(size)
+        return Xr
+
+
+
+def ksparse(x, nb_active=10):
+    mask = torch.ones(x.size())
+    for i, xi in enumerate(x.data.tolist()):
+        inds = np.argsort(xi)
+        inds = inds[::-1]
+        inds = inds[nb_active:]
+        if len(inds):
+            inds = np.array(inds)
+            inds = torch.from_numpy(inds).long()
+            mask[i][inds] = 0
+    return x * (mask).float().to(x.device)
+
+
+class ConvAE(nn.Module):
+
+    def __init__(self, w=32, h=32, c=1, nb_filters=64, spatial=True, channel=True, channel_stride=4):
+        super().__init__()
+        self.spatial = spatial
+        self.channel = channel
+        self.channel_stride = channel_stride
+        self.encode = nn.Sequential(
+            nn.Conv2d(c, nb_filters, 5, 1, 0),
+            nn.ReLU(True),
+            nn.Conv2d(nb_filters, nb_filters, 5, 1, 0),
+            nn.ReLU(True),
+            nn.Conv2d(nb_filters, nb_filters, 5, 1, 0),
+        )
+        self.decode = nn.Sequential(
+            nn.ConvTranspose2d(nb_filters, c, 13, 1, 0),
+            nn.Sigmoid()
+        )
+        self.apply(_weights_init)
+
+    def forward(self, X):
+        size = X.size()
+        h = self.encode(X)
+        h = self.sparsify(h)
+        Xr = self.decode(h)
+        return Xr
+    
+    def sparsify(self, h):
+        if self.spatial:
+            h = spatial_sparsity(h)
+        if self.channel:
+            h = strided_channel_sparsity(h, stride=self.channel_stride)
+        return h
+
+class SimpleConvAE(nn.Module):
+
+    def __init__(self, w=32, h=32, c=1, nb_filters=64, spatial=True, channel=True, channel_stride=4):
+        super().__init__()
+        self.spatial = spatial
+        self.channel = channel
+        self.channel_stride = channel_stride
+        self.encode = nn.Sequential(
+            nn.Conv2d(c, nb_filters, 13, 1, 0),
+            nn.ReLU(True),
+        )
+        self.decode = nn.Sequential(
+            nn.ConvTranspose2d(nb_filters, c, 13, 1, 0),
+            nn.Sigmoid()
+        )
+        self.apply(_weights_init)
+
+    def forward(self, X):
+        size = X.size()
+        h = self.encode(X)
+        h = self.sparsify(h)
+        Xr = self.decode(h)
+        return Xr
+    
+    def sparsify(self, h):
+        if self.spatial:
+            h = spatial_sparsity(h)
+        if self.channel:
+            h = strided_channel_sparsity(h, stride=self.channel_stride)
+        return h
+
+class DeepConvAE(nn.Module):
+
+    def __init__(self, w=32, h=32, c=1, nb_filters=64, nb_layers=3, spatial=True, channel=True, channel_stride=4):
+        super().__init__()
+        self.spatial = spatial
+        self.channel = channel
+        self.channel_stride = channel_stride
+
+        layers = [
+            nn.Conv2d(c, nb_filters, 5, 1, 0),
+            nn.ReLU(True),
+        ]
+        for _ in range(nb_layers - 1):
+            layers.extend([
+                nn.Conv2d(nb_filters, nb_filters, 5, 1, 0),
+                nn.ReLU(True),
+            ])
+        self.encode = nn.Sequential(*layers)
+        layers = []
+        for _ in range(nb_layers - 1):
+            layers.extend([
+                nn.ConvTranspose2d(nb_filters, nb_filters, 5, 1, 0),
+                nn.ReLU(True),
+            ])
+        layers.extend([
+            nn.ConvTranspose2d(nb_filters, c, 5, 1, 0),
+            nn.Sigmoid()
+        ])
+        self.decode = nn.Sequential(*layers)
+        self.apply(_weights_init)
+
+    def forward(self, X):
+        size = X.size()
+        h = self.encode(X)
+        h = self.sparsify(h)
+        Xr = self.decode(h)
+        return Xr
+    
+    def sparsify(self, h):
+        if self.spatial:
+            h = spatial_sparsity(h)
+        if self.channel:
+            h = strided_channel_sparsity(h, stride=self.channel_stride)
+        return h
+
+
+def spatial_sparsity(x):
+    maxes = x.amax(dim=(2,3), keepdims=True)
+    return x * equals(x, maxes)
+
+def equals(x, y, eps=1e-8):
+    return torch.abs(x-y) <= eps
+
+def strided_channel_sparsity(x, stride=1):
+    B, F = x.shape[0:2]
+    h, w = x.shape[2:]
+    x_ = x.view(B, F, h // stride, stride, w // stride, stride)
+    mask = equals(x_, x_.amax(axis=(1, 3, 5), keepdims=True))
+    mask = mask.view(x.shape).float()
+    return x * mask
+
+
+def _weights_init(m):
+    if hasattr(m, 'weight'):
+        xavier_uniform(m.weight.data)
+        if m.bias is not None:
+            m.bias.data.fill_(0)

test.py

@@ -0,0 +1,21 @@
+import torch
+import numpy as np
+from machinedesign.autoencoder.interface import load
+from keras.models import Model
+torch.use_deterministic_algorithms(True)
+model = torch.load("mnist_deepconvae/model.th")
+model_keras = load("/home/mehdi/work/code/out_of_class/ae/mnist")
+print(model_keras.layers[8])
+
+m = Model(model_keras.inputs, model_keras.layers[8].output)
+X = torch.rand(1,1,28,28)
+with torch.no_grad():
+    # X1 = model.sparsify(model.encode(X))
+    X1 = model(X)
+X2 = model_keras.predict(X)
+X2 = torch.from_numpy(X2)
+print(torch.abs(X1-X2).sum())
+# for i in range(128):
+    # print(i, torch.abs(X1[0,i]-X2[0,i]).sum())
+    # print(X1[0,i, 0, :])
+    # print(X2[0,i,0, :])

viz.py

@@ -0,0 +1,204 @@
+"""
+This module contains common visualization functions
+used to report results of the models.
+"""
+
+from functools import partial
+import numpy as np
+
+
+def horiz_merge(left, right):
+    """
+    merges two images, left and right horizontally to obtain
+    a bigger image containing both.
+
+    Parameters
+    ---------
+    left: 2D or 3D numpy array
+        left image.
+        2D for grayscale.
+        3D for color.
+    right : numpy array array
+        right image.
+        2D for grayscale
+        3D for color.
+
+    Returns
+    -------
+
+    numpy array (2D or 3D depending on left and right)
+    """
+    assert left.shape[0] == right.shape[0]
+    assert left.shape[2:] == right.shape[2:]
+    shape = (left.shape[0], left.shape[1] + right.shape[1],) + left.shape[2:]
+    im_merge = np.zeros(shape)
+    im_merge[:, 0:left.shape[1]] = left
+    im_merge[:, left.shape[1]:] = right
+    return im_merge
+
+def vert_merge(top, bottom):
+    """
+    merges two images, top and bottom vertically to obtain
+    a bigger image containing both.
+
+    Parameters
+    ---------
+    top: 2D or 3D numpy array
+        top image.
+        2D for grayscale.
+        3D for color.
+    bottom : numpy array array
+        bottom image.
+        2D for grayscale
+        3D for color.
+
+    Returns
+    -------
+
+    numpy array (2D or 3D depending on left and right)
+    """
+    im = horiz_merge(top, bottom)
+    if len(im.shape) == 2:
+        im = im.transpose((1, 0))
+    elif len(im.shape) == 3:
+        im = im.transpose((1, 0, 2))
+    return im
+
+
+def grid_of_images(M, border=0, bordercolor=[0.0, 0.0, 0.0], shape=None, normalize=False):
+    """
+    Draw a grid of images from M
+    The order in the grid which corresponds to the order in M
+    is starting from top to bottom then left to right.
+
+    Parameters
+    ----------
+
+    M : numpy array
+        if 3D, convert it to 4D, the shape will be interpreted as (nb_images, h, w) and converted to (nb_images, 1, h, w).
+        if 4D, consider it as colored or grayscale
+            - if the shape is (nb_images, nb_colors, h, w), it is converted to (nb_images, h, w, nb_colors)
+            - otherwise, if it already (nb_images, h, w, nb_colors), use it as it is.
+            - nb_colors can be 1 (grayscale) or 3 (colors).
+    border: int
+        thickness of border(default=0)
+    shape: tuple (nb_cols, nb_rows)
+        shape of the grid
+        by default make a square shape
+        (in that case, it is possible that not all images from M will be part of the grid).
+    normalize: bool(default=False)
+        whether to normalize the pixel values of each image independently
+        by min and max. if False, clip the values of pixels to 0 and 1
+        without normalizing.
+
+    Returns
+    -------
+
+    3D numpy array of shape (h, w, 3)
+    (with a color channel regardless of whether the original images were grayscale or colored)
+    """
+    if len(M.shape) == 3:
+        M = M[:, :, :, np.newaxis]
+    if M.shape[-1] not in (1, 3):
+        M = M.transpose((0, 2, 3, 1))
+    if M.shape[-1] == 1:
+        M = np.ones((1, 1, 1, 3)) * M
+    bordercolor = np.array(bordercolor)[None, None, :]
+    numimages = len(M)
+    M = M.copy()
+
+    if normalize:
+        for i in range(M.shape[0]):
+            M[i] -= M[i].flatten().min()
+            M[i] /= M[i].flatten().max()
+    else:
+        M = np.clip(M, 0, 1)
+    height, width, color = M[0].shape
+    assert color == 3, 'Nb of color channels are {}'.format(color)
+    if shape is None:
+        n0 = np.int(np.ceil(np.sqrt(numimages)))
+        n1 = np.int(np.ceil(np.sqrt(numimages)))
+    else:
+        n0 = shape[0]
+        n1 = shape[1]
+
+    im = np.array(bordercolor) * np.ones(
+        ((height + border) * n1 + border, (width + border) * n0 + border, 1), dtype='<f8')
+    # shape = (n0, n1)
+    # j corresponds to rows in the grid, n1 should correspond to nb of rows
+    # i corresponds to columns in the grid, n0 should correspond to nb of cols
+    # M should be such that the first n1 examples correspond to row 1, 
+    # next n1 examples correspond to row 2, etc. that is, M first axis
+    # can be reshaped to (n1, n0)
+    for i in range(n0):
+        for j in range(n1):
+            if i * n1 + j < numimages:
+                im[j * (height + border) + border:(j + 1) * (height + border) + border,
+                   i * (width + border) + border:(i + 1) * (width + border) + border, :] = np.concatenate((
+                       np.concatenate((M[i * n1 + j, :, :, :],
+                                       bordercolor * np.ones((height, border, 3), dtype=float)), 1),
+                       bordercolor * np.ones((border, width + border, 3), dtype=float)
+                   ), 0)
+    return im
+
+grid_of_images_default = partial(grid_of_images, border=1, bordercolor=(0.3, 0, 0))
+
+
+def reshape_to_images(x, input_shape=None):
+    """
+    a function that takes a numpy array and try to
+    reshape it to an array of images that would
+    be compatible with the function grid_of_images.
+    Two cases are considered.
+
+    if x is a 2D numpy array, it uses input_shape:
+        - x can either be (nb_examples, nb_features) or (nb_features, nb_examples)
+        - nb_features should be prod(input_shape)
+        - the nb_features dim is then expanded to have :
+            (nb_examples, h, w, nb_channels), sorted input_shape shoud
+            be  (h, w, nb_channels).
+
+    if x is a 4D numpy array:
+        - if the first tensor dim is 1 or 3 like e.g. (1, a, b, c), then assume it is 
+          color channel and transform to (a, 1, b, c)
+        - if the second tensor dim is 1 or 3, leave x it as it is
+        - if the third tensor dim is 1 or 3, like e.g. (a, b, 1, c), then assume it is
+          color channel and transform to (c, 1, a, b)
+        - if the fourth tensor dim is 1 or 3, like e.g. (a, b, c, 1), then assume it is
+          color channel and transform to (c, 1, a, b)
+    Parameters
+    ----------
+
+    x : numpy array
+        input to be reshape
+    input_shape : tuple needed only when x is 2D numpy array
+    """
+    if len(x.shape) == 2:
+        assert input_shape is not None
+        if x.shape[0] == np.prod(input_shape):
+            x = x.T
+            x = x.reshape((x.shape[0],) + input_shape)
+            x = x.transpose((0, 2, 3, 1))
+            return x
+        elif x.shape[1] == np.prod(input_shape):
+            x = x.reshape((x.shape[0],) + input_shape)
+            x = x.transpose((0, 2, 3, 1))
+            return x
+        else:
+            raise ValueError('Cant recognize this shape : {}'.format(x.shape))
+    elif len(x.shape) == 4:
+        if x.shape[0] in (1, 3):
+            x = x.transpose((1, 0, 2, 3))
+            return x
+        elif x.shape[1] in (1, 3):
+            return x
+        elif x.shape[2] in (1, 3):
+             x = x.transpose((3, 2, 0, 1))
+             return x
+        elif x.shape[3] in (1, 3):
+            x = x.transpose((2, 3, 0, 1))
+            return x
+        else:
+            raise ValueError('Cant recognize a shape of size : {}'.format(len(x.shape)))
+    else:
+        raise ValueError('Cant recognize a shape of size : {}'.format(len(x.shape)))