Merge pull request #3 from ritheshkumar95/tristan/miniimagenet

Mini-Imagenet experiments

Author: ritheshkumar95

.gitignore

@@ -107,6 +107,7 @@ run.sh
 # Logs & Saves
 logs/
 saves/
+models/
 
 # Slurm
 *.out

datasets.py

@@ -0,0 +1,117 @@
+import os
+import csv
+import torch.utils.data as data
+from PIL import Image
+
+def pil_loader(path):
+    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
+    # Borrowed from https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py
+    with open(path, 'rb') as f:
+        img = Image.open(f)
+        return img.convert('RGB')
+
+class MiniImagenet(data.Dataset):
+
+    base_folder = '/data/lisa/data/miniimagenet'
+    filename = 'miniimagenet.zip'
+    splits = {
+        'train': 'train.csv',
+        'valid': 'val.csv',
+        'test': 'test.csv'
+    }
+
+    def __init__(self, root, train=False, valid=False, test=False,
+                 transform=None, target_transform=None, download=False):
+        super(MiniImagenet, self).__init__()
+        self.root = root
+        self.train = train
+        self.valid = valid
+        self.test = test
+        self.transform = transform
+        self.target_transform = target_transform
+
+        if not (((train ^ valid ^ test) ^ (train & valid & test))):
+            raise ValueError('One and only one of `train`, `valid` or `test` '
+                'must be True (train={0}, valid={1}, test={2}).'.format(train,
+                valid, test))
+
+        self.image_folder = os.path.join(os.path.expanduser(root), 'images')
+        if train:
+            split = self.splits['train']
+        elif valid:
+            split = self.splits['valid']
+        elif test:
+            split = self.splits['test']
+        else:
+            raise ValueError('Unknown split.')
+        self.split_filename = os.path.join(os.path.expanduser(root), split)
+        if download:
+            self.download()
+        if not self._check_exists():
+            raise RuntimeError('Dataset not found. You can use `download=True` '
+                               'to download it')
+
+        # Extract filenames and labels
+        self._data = []
+        with open(self.split_filename, 'r') as f:
+            reader = csv.reader(f)
+            next(reader) # Skip the header
+            for line in reader:
+                self._data.append(tuple(line))
+        self._fit_label_encoding()
+
+    def __getitem__(self, index):
+        filename, label = self._data[index]
+        image = pil_loader(os.path.join(self.image_folder, filename))
+        label = self._label_encoder[label]
+        if self.transform is not None:
+            image = self.transform(image)
+        if self.target_transform is not None:
+            label = self.target_transform(label)
+
+        return image, label
+
+    def _fit_label_encoding(self):
+        _, labels = zip(*self._data)
+        unique_labels = set(labels)
+        self._label_encoder = dict((label, idx)
+            for (idx, label) in enumerate(unique_labels))
+
+    def _check_exists(self):
+        return (os.path.exists(self.image_folder) 
+            and os.path.exists(self.split_filename))
+
+    def download(self):
+        from shutil import copyfile
+        from zipfile import ZipFile
+
+        # If the image folder already exists, break
+        if self._check_exists():
+            return True
+
+        # Create folder if it does not exist
+        root = os.path.expanduser(self.root)
+        if not os.path.exists(root):
+            os.makedirs(root)
+
+        # Copy the file to root
+        path_source = os.path.join(self.base_folder, self.filename)
+        path_dest = os.path.join(root, self.filename)
+        print('Copy file `{0}` to `{1}`...'.format(path_source, path_dest))
+        copyfile(path_source, path_dest)
+
+        # Extract the dataset
+        print('Extract files from `{0}`...'.format(path_dest))
+        with ZipFile(path_dest, 'r') as f:
+            f.extractall(root)
+
+        # Copy CSV files
+        for split in self.splits:
+            path_source = os.path.join(self.base_folder, self.splits[split])
+            path_dest = os.path.join(root, self.splits[split])
+            print('Copy file `{0}` to `{1}`...'.format(path_source, path_dest))
+            copyfile(path_source, path_dest)
+        print('Done!')
+
+    def __len__(self):
+        return len(self._data)

miniimagenet_pixelcnn_prior.py

@@ -0,0 +1,171 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import json
+from torchvision import transforms
+from torchvision.utils import save_image, make_grid
+
+from modules import AutoEncoder, GatedPixelCNN
+from datasets import MiniImagenet
+
+from tensorboardX import SummaryWriter
+
+def train(data_loader, model, prior, optimizer, args, writer):
+    for images, labels in data_loader:
+        with torch.no_grad():
+            images = images.to(args.device)
+            latents, _ = model.encode(images)
+            latents = latents.detach()
+
+        labels = labels.to(args.device)
+        logits = prior(latents, labels)
+        logits = logits.permute(0, 2, 3, 1).contiguous()
+
+        optimizer.zero_grad()
+        loss = F.cross_entropy(logits.view(-1, args.k),
+                               latents.view(-1))
+        loss.backward()
+
+        # Logs
+        writer.add_scalar('loss/train', loss.item(), args.steps)
+
+        optimizer.step()
+        args.steps += 1
+
+def test(data_loader, model, prior, args, writer):
+    with torch.no_grad():
+        loss = 0.
+        for images, labels in data_loader:
+            images = images.to(args.device)
+            labels = labels.to(args.device)
+
+            latents, _ = model.encode(images)
+            latents = latents.detach()
+            logits = prior(latents, labels)
+            logits = logits.permute(0, 2, 3, 1).contiguous()
+            loss += F.cross_entropy(logits.view(-1, args.k),
+                                    latents.view(-1))
+
+        loss /= len(data_loader)
+
+    # Logs
+    writer.add_scalar('loss/valid', loss.item(), args.steps)
+
+    return loss.item()
+
+def main(args):
+    writer = SummaryWriter('./logs/{0}'.format(args.output_folder))
+    save_filename = './models/{0}/prior.pt'.format(args.output_folder)
+
+    transform = transforms.Compose([
+        transforms.RandomResizedCrop(128),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+    ])
+
+    # Define the train, valid & test datasets
+    train_dataset = MiniImagenet(args.data_folder, train=True,
+        download=True, transform=transform)
+    valid_dataset = MiniImagenet(args.data_folder, valid=True,
+        download=True, transform=transform)
+    test_dataset = MiniImagenet(args.data_folder, test=True,
+        download=True, transform=transform)
+    # Define the data loaders
+    train_loader = torch.utils.data.DataLoader(train_dataset,
+        batch_size=args.batch_size, shuffle=False,
+        num_workers=args.num_workers, pin_memory=True)
+    valid_loader = torch.utils.data.DataLoader(valid_dataset,
+        batch_size=args.batch_size, shuffle=False, drop_last=True,
+        num_workers=args.num_workers, pin_memory=True)
+    test_loader = torch.utils.data.DataLoader(test_dataset,
+        batch_size=16, shuffle=True)
+
+    # Save the label encoder
+    with open('./models/{0}/labels.json'.format(args.output_folder), 'w') as f:
+        json.dump(train_dataset._label_encoder, f)
+
+    # Fixed images for Tensorboard
+    fixed_images, _ = next(iter(test_loader))
+    fixed_grid = make_grid(fixed_images, nrow=8, range=(-1, 1), normalize=True)
+    writer.add_image('original', fixed_grid, 0)
+
+    model = AutoEncoder(3, args.hidden_size_vae, args.k).to(args.device)
+    with open(args.model, 'rb') as f:
+        state_dict = torch.load(f)
+        model.load_state_dict(state_dict)
+    model.eval()
+
+    prior = GatedPixelCNN(args.k, args.hidden_size_prior,
+        args.num_layers, n_classes=len(train_dataset._label_encoder)).to(args.device)
+    optimizer = torch.optim.Adam(prior.parameters(), lr=args.lr)
+
+    best_loss = -1.
+    for epoch in range(args.num_epochs):
+        train(train_loader, model, prior, optimizer, args, writer)
+        # The validation loss is not properly computed since
+        # the classes in the train and valid splits of Mini-Imagenet
+        # do not overlap.
+        loss = test(valid_loader, model, prior, args, writer)
+
+        if (epoch == 0) or (loss < best_loss):
+            best_loss = loss
+            with open(save_filename, 'wb') as f:
+                torch.save(prior.state_dict(), f)
+
+if __name__ == '__main__':
+    import argparse
+    import os
+    import multiprocessing as mp
+
+    parser = argparse.ArgumentParser(description='PixelCNN Prior for VQ-VAE')
+
+    # General
+    parser.add_argument('--data-folder', type=str,
+        help='name of the data folder')
+    parser.add_argument('--model', type=str,
+        help='filename containing the model')
+
+    # Latent space
+    parser.add_argument('--hidden-size-vae', type=int, default=256,
+        help='size of the latent vectors (default: 256)')
+    parser.add_argument('--hidden-size-prior', type=int, default=64,
+        help='hidden size for the PixelCNN prior (default: 64)')
+    parser.add_argument('--k', type=int, default=512,
+        help='number of latent vectors (default: 512)')
+    parser.add_argument('--num-layers', type=int, default=15,
+        help='number of layers for the PixelCNN prior (default: 15)')
+
+    # Optimization
+    parser.add_argument('--batch-size', type=int, default=128,
+        help='batch size (default: 128)')
+    parser.add_argument('--num-epochs', type=int, default=100,
+        help='number of epochs (default: 100)')
+    parser.add_argument('--lr', type=float, default=3e-4,
+        help='learning rate for Adam optimizer (default: 3e-4)')
+
+    # Miscellaneous
+    parser.add_argument('--output-folder', type=str, default='prior',
+        help='name of the output folder (default: prior)')
+    parser.add_argument('--num-workers', type=int, default=mp.cpu_count() - 1,
+        help='number of workers for trajectories sampling (default: {0})'.format(mp.cpu_count() - 1))
+    parser.add_argument('--device', type=str, default='cpu',
+        help='set the device (cpu or cuda, default: cpu)')
+
+    args = parser.parse_args()
+
+    # Create logs and models folder if they don't exist
+    if not os.path.exists('./logs'):
+        os.makedirs('./logs')
+    if not os.path.exists('./models'):
+        os.makedirs('./models')
+    # Device
+    args.device = torch.device(args.device
+        if torch.cuda.is_available() else 'cpu')
+    # Slurm
+    if 'SLURM_JOB_ID' in os.environ:
+        args.output_folder += '-{0}'.format(os.environ['SLURM_JOB_ID'])
+    if not os.path.exists('./models/{0}'.format(args.output_folder)):
+        os.makedirs('./models/{0}'.format(args.output_folder))
+    args.steps = 0
+
+    main(args)