main.py

from IPython import embed
import argparse
from operator import itemgetter
import os
import time

from PIL import Image
from tensorboardX import SummaryWriter
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim import lr_scheduler
from torchvision import datasets, transforms, utils
import tqdm

from masking import *
from model import *
from ours import *
from utils import *

parser = argparse.ArgumentParser()
# data I/O
parser.add_argument('-i', '--data_dir', type=str,
                    default='data', help='Location for the dataset')
parser.add_argument('-o', '--save_dir', type=str, default='models',
                    help='Location for parameter checkpoints and samples')
parser.add_argument('-d', '--dataset', type=str,
                    default='cifar', help='Can be either cifar|mnist|celeba')
parser.add_argument('-p', '--print_every', type=int, default=20,
                    help='how many iterations between print statements')
parser.add_argument('-t', '--save_interval', type=int, default=20,
                    help='Every how many epochs to write checkpoint?')
parser.add_argument('-ts', '--sample_interval', type=int, default=4,
                    help='Every how many epochs to write samples?')
parser.add_argument('-r', '--load_params', type=str, default=None,
                    help='Restore training from previous model checkpoint?')
parser.add_argument('-rd', '--run_dir', type=str, default=None,
                    help="Optionally specify run directory. One will be generated otherwise."
                         "Use to save log files in a particular place")
parser.add_argument('--exp_name', type=str, default=None)
parser.add_argument('-ID', '--exp_id', type=int, default=0)
parser.add_argument('--ours', action='store_true')
# pixelcnn++ and our model
parser.add_argument('-q', '--nr_resnet', type=int, default=5,
                    help='Number of residual blocks per stage of the model')
parser.add_argument('-n', '--nr_filters', type=int, default=160,
                    help='Number of filters to use across the model. Higher = larger model.')
parser.add_argument('-m', '--nr_logistic_mix', type=int, default=10,
                    help='Number of logistic components in the mixture. Higher = more flexible model')
parser.add_argument('-l', '--lr', type=float,
                    default=0.0002, help='Base learning rate')
parser.add_argument('-e', '--lr_decay', type=float, default=0.999995,
                    help='Learning rate decay, applied every step of the optimization')
parser.add_argument('-wd', '--weight_decay', type=float,
                    default=0, help='Weight decay during optimization')
parser.add_argument('-c', '--clip', type=float, default=-1, help='Gradient norms clipped to this value')
parser.add_argument('-b', '--batch_size', type=int, default=64,
                    help='Batch size during training per GPU')
parser.add_argument('-x', '--max_epochs', type=int,
                    default=5000, help='How many epochs to run in total?')
parser.add_argument('-s', '--seed', type=int, default=1,
                    help='Random seed to use')
# our model
parser.add_argument('-k', '--kernel_size', type=int, default=5,
                    help='Size of conv kernels')
parser.add_argument('-md', '--max_dilation', type=int, default=2,
                    help='Dilation in downsize stream')
parser.add_argument('-dp', '--dropout_prob', type=float, default=0.5,
                    help='Dropout prob used with nn.Dropout2d in gated resnet layers. '
                         'Argument only used if --ours is provided. Set to 0 to disable '
                         'dropout entirely.')
parser.add_argument('-nm', '--normalization', type=str, default='weight_norm',
                    choices=["none", "weight_norm", "instance_norm", "instance_norm_affine",
                             "order_rescale", "pono"])
parser.add_argument('-af', '--accum_freq', type=int, default=1,
                    help='Batches per optimization step. Used for gradient accumulation')
parser.add_argument('--two_stream', action="store_true", help="Enable two stream model")
parser.add_argument('--order', type=str, choices=["raster_scan", "s_curve", "hilbert", "gilbert2d"],
                    help="Autoregressive generation order")
parser.add_argument('--randomize_order', action="store_true", help="Randomize between 8 variants of the "
                    "pixel generation order.")
parser.add_argument('--mode', type=str, choices=["train", "sample", "test"],
                    default="train")
parser.add_argument('--no_bias', action="store_true", help="Disable learnable bias for all convolutions")
parser.add_argument('--minimize_bpd', action="store_true", help="Minimize bpd, scaling loss down by number of dimension")
parser.add_argument('--resize_sizes', type=int, nargs="*")
parser.add_argument('--resize_probs', type=float, nargs="*")

args = parser.parse_args()
# assert args.normalization != "weight_norm", "Weight normalization manually disabled in layers.py"


# Set seed for reproducibility
torch.manual_seed(args.seed)
np.random.seed(args.seed)


# Create run directory
if args.run_dir:
    run_dir = args.run_dir
else:
    _name = "{:05d}_{}_lr{:.5f}_bs{}_gc{}_k{}_md{}".format(
        args.exp_id, args.dataset, args.lr, args.batch_size, args.clip, args.kernel_size, args.max_dilation)
    if args.normalization != "none":
        _name = f"{_name}_{args.normalization}"
    if args.exp_name:
        _name = f"{_name}+{args.exp_name}"
    run_dir = os.path.join("runs", _name)
    if args.mode == "train":
        os.makedirs(run_dir, exist_ok=False)
assert os.path.exists(run_dir), "Did not find run directory, check --run_dir argument"


# Log arguments
logger = configure_logger(os.path.join(run_dir, f"{args.mode}.log"))
logger.info("Run directory: %s", run_dir)
logger.info("Arguments: %s", args)
for k, v in vars(args).items():
    logger.info(f"  {k}: {v}")


# Create data loaders
sample_batch_size = 25
dataset_obs = {
    'mnist': (1, 28, 28),
    'cifar': (3, 32, 32),
    'celebahq': (3, 256, 256)
}[args.dataset]
input_channels = dataset_obs[0]
data_loader_kwargs = {'num_workers':1, 'pin_memory':True, 'drop_last':True, 'batch_size':args.batch_size}
if args.resize_sizes:
    if not args.resize_probs:
        args.resize_probs = [1. / len(args.resize_sizes)] * len(args.resize_sizes)
    assert len(args.resize_probs) == len(args.resize_sizes)
    assert sum(args.resize_probs) == 1
    resized_obses = [(input_channels, s, s) for s in args.resize_sizes]
else:
    args.resize_sizes = [dataset_obs[1]]
    args.resize_probs = [1.]
    resized_obses = [dataset_obs]

def obs2str(obs):
    return 'x'.join(map(str, obs))

def random_resized_obs():
    idx = np.arange(len(resized_obses))
    obs_i = np.random.choice(idx, p=args.resize_probs)
    return resized_obses[int(obs_i)]

def get_resize_collate_fn(obs, default_collate=torch.utils.data.dataloader.default_collate):
    if obs == dataset_obs:
        return default_collate

    def resize_collate_fn(batch):
        X, y = default_collate(batch)
        X = torch.nn.functional.interpolate(X, size=obs[1:], mode="bilinear")
        return [X, y]
    return resize_collate_fn

def random_resize_collate(batch):
    X, y = torch.utils.data.dataloader.default_collate(batch)
    obs = random_resized_obs()
    if obs != dataset_obs:
        X = torch.nn.functional.interpolate(X, size=obs[1:], mode="bilinear")
    return [X, y]

# Create data loaders
if 'mnist' in args.dataset :
    rescaling = lambda x : (x - .5) * 2.  # rescale [0, 1] images into [-1, 1] range
    rescaling_inv = lambda x : .5 * x + .5
    ds_transforms = transforms.Compose([transforms.ToTensor(), rescaling])

    train_loader = torch.utils.data.DataLoader(datasets.MNIST(args.data_dir, download=True,
        train=True, transform=ds_transforms), shuffle=True, collate_fn=random_resize_collate, **data_loader_kwargs)
    test_loader_by_obs = {
        obs: torch.utils.data.DataLoader(datasets.MNIST(args.data_dir, train=False,
            transform=ds_transforms), collate_fn=get_resize_collate_fn(obs), **data_loader_kwargs)
        for obs in resized_obses
    }
elif 'cifar' in args.dataset :
    rescaling = lambda x : (x - .5) * 2.  # rescale [0, 1] images into [-1, 1] range
    rescaling_inv = lambda x : .5 * x + .5
    ds_transforms = transforms.Compose([transforms.ToTensor(), rescaling])

    train_loader = torch.utils.data.DataLoader(datasets.CIFAR10(args.data_dir, train=True, 
        download=True, transform=ds_transforms), shuffle=True, collate_fn=random_resize_collate, **data_loader_kwargs)
    test_loader_by_obs = {
        obs: torch.utils.data.DataLoader(datasets.CIFAR10(args.data_dir, train=False,
            transform=ds_transforms), collate_fn=get_resize_collate_fn(obs), **data_loader_kwargs)
        for obs in resized_obses
    }
elif 'celebahq' in args.dataset :
    rescaling = lambda x : (2. / 255) * x - 1.  # rescale uint8 images into [-1, 1] range
    rescaling_inv = lambda x : (255. / 2) * (x + 1.)

    # NOTE: Random resizing of images during training is not supported for CelebA-HQ. Will use 256x256 resolution.
    from celeba_data import get_celeba_dataloader
    del data_loader_kwargs["num_workers"]
    train_loader = get_celeba_dataloader(args.data_dir, "train",
                                         collate_fn=itemgetter(0),
                                         batch_transform=rescaling,
                                         **data_loader_kwargs)
    test_loader_by_obs = {
        obs: get_celeba_dataloader(args.data_dir, "validation",
                                   collate_fn=get_resize_collate_fn(obs, itemgetter(0)),
                                   batch_transform=rescaling,
                                   **data_loader_kwargs)
        for obs in resized_obses
    }
else :
    raise Exception('{} dataset not in {mnist, cifar10}'.format(args.dataset))


# Select loss functions
if 'mnist' in args.dataset :
    # Losses for 1-channel images
    loss_op = discretized_mix_logistic_loss_1d
    loss_op_averaged = discretized_mix_logistic_loss_1d_averaged
    sample_op = lambda x : sample_from_discretized_mix_logistic_1d(x, args.nr_logistic_mix)
else:
    # Losses for 3-channel images
    loss_op = discretized_mix_logistic_loss
    loss_op_averaged = discretized_mix_logistic_loss_averaged
    sample_op = lambda x : sample_from_discretized_mix_logistic(x, args.nr_logistic_mix)


# Construct model
if args.ours:
    logger.info("Constructing our model")

    if args.normalization == "instance_norm":
        raise NotImplementedError("Causal instance norm not implemented")
        # norm_op = lambda num_channels: nn.InstanceNorm2d(num_channels)
    elif args.normalization == "instance_norm_affine":
        raise NotImplementedError("Causal instance norm not implemented")
        # norm_op = lambda num_channels: nn.InstanceNorm2d(num_channels, affine=True)
    elif args.normalization == "order_rescale":
        norm_op = lambda num_channels: OrderRescale()
    elif args.normalization == "pono":
        norm_op = lambda num_channels: PONO()
    else:
        norm_op = None

    assert not args.two_stream, "--two_stream cannot be used with --ours"
    model = OurPixelCNN(
                nr_resnet=args.nr_resnet,
                nr_filters=args.nr_filters, 
                input_channels=input_channels,
                nr_logistic_mix=args.nr_logistic_mix,
                kernel_size=(args.kernel_size, args.kernel_size),
                max_dilation=args.max_dilation,
                weight_norm=(args.normalization == "weight_norm"),
                feature_norm_op=norm_op,
                dropout_prob=args.dropout_prob,
                conv_bias=(not args.no_bias))

    all_generation_idx_by_obs = {}
    all_masks_by_obs = {}
    for obs in resized_obses:
        # Get generation orders
        base_generation_idx = get_generation_order_idx(args.order, obs[1], obs[2])
        if args.randomize_order:
            all_generation_idx = augment_orders(base_generation_idx, obs)
        else:
            all_generation_idx = [base_generation_idx]
        if args.mode == "train":
            plot_orders(all_generation_idx, obs, size=5, plot_rows=min(len(all_generation_idx), 4),
                        out_path=os.path.join(run_dir, f"orderings_obs{obs2str(obs)}.png"))
        all_generation_idx_by_obs[obs] = all_generation_idx

        # Make masks and plot
        all_masks = []
        for i, generation_idx in enumerate(all_generation_idx):
            masks = get_masks(generation_idx, obs[1], obs[2], args.kernel_size, args.max_dilation,
                            run_dir, plot_suffix=f"obs{obs2str(obs)}_order{i}", plot=False)#(args.mode == "train"))
            logger.info(f"Mask shapes: {masks[0].shape}, {masks[1].shape}, {masks[2].shape}")
            all_masks.append(masks)
        all_masks_by_obs[obs] = all_masks
else:
    logger.info("Constructing original PixelCNN++")
    model = PixelCNN(nr_resnet=args.nr_resnet, nr_filters=args.nr_filters, 
                input_channels=input_channels, nr_logistic_mix=args.nr_logistic_mix)

    assert not args.randomize_order
    all_generation_idx_by_obs = {}
    all_masks_by_obs = {}
    for obs in resized_obses:
        all_generation_idx_by_obs[obs] = [get_generation_order_idx("raster_scan", obs[1], obs[2])]
        all_masks_by_obs[obs] = [(None, None, None)]
model = nn.DataParallel(model)
model = model.cuda()


# Load model parameters from checkpoint
if args.load_params:
    # TODO: Restore optimizer
    if os.path.exists(args.load_params):
        load_params = args.load_params
    else:
        load_params = os.path.join(run_dir, args.load_params)
    checkpoint_epochs = load_part_of_model(load_params, model=model.module, optimizer=None)
    logger.info(f"Model parameters loaded from {load_params}, from after {checkpoint_epochs} training epochs")
else:
    checkpoint_epochs = -1


# Create optimizer
# NOTE: PixelCNN++ TF repo uses betas=(0.95, 0.9995), different than PyTorch defaults
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
scheduler = lr_scheduler.StepLR(optimizer, step_size=1, gamma=args.lr_decay)


def test(model, all_masks, test_loader, epoch="N/A", progress_bar=True):
    logger.info(f"Testing with ensemble of {len(all_masks)} orderings")
    test_loss = 0.
    for batch_idx, (input,_) in enumerate(tqdm.tqdm(test_loader,
                                                    desc=f"Test after epoch {epoch}",
                                                    disable=(not progress_bar))):
        input = input.cuda(non_blocking=True)
        input_var = Variable(input)

        #mask_init, mask_undilated, mask_dilated = all_masks[0]
        #output = model(input_var, mask_init=mask_init, mask_undilated=mask_undilated, mask_dilated=mask_dilated)
        #loss = loss_op(input_var, output)

        # Average likelihoods over multiple orderings
        outputs = []
        for mask_init, mask_undilated, mask_dilated in all_masks:
            output = model(input_var, mask_init=mask_init, mask_undilated=mask_undilated, mask_dilated=mask_dilated)
            outputs.append(output)
        loss = loss_op_averaged(input_var, outputs)

        test_loss += loss.item()
        del loss, output

    # FIXME: for final evaluation, don't use batch_idx * args.batch_size -- this slightly overestimates
    # the number of dims (10016 * prod(obs) * log(2) for mnist) since the last iteration might have fewer than
    # args.batch_size images. Leaving this code the same for now to allow comparison between training runs.
    deno = batch_idx * args.batch_size * np.prod(obs) * np.log(2.)
    assert deno > 0, embed()
    test_bpd = test_loss / deno
    return test_bpd


def sample(model, generation_idx, mask_init, mask_undilated, mask_dilated):
    model.eval()
    data = torch.zeros(sample_batch_size, obs[0], obs[1], obs[2])
    data = data.cuda()
    for num_px_sampled, (i, j) in enumerate(tqdm.tqdm(generation_idx, desc="Sampling pixels")):
        data_v = Variable(data)
        out = model(data_v, sample=True, mask_init=mask_init, mask_undilated=mask_undilated, mask_dilated=mask_dilated)
        out_sample = sample_op(out)
        data[:, :, i, j] = out_sample.data[:, :, i, j]
    return data


if args.mode == "train":
    logger.info("starting training")
    writer = SummaryWriter(log_dir=run_dir)
    global_step = 0
    min_train_bpd = 1e12
    min_test_bpd_by_obs = {obs: 1e12 for obs in resized_obses}
    for epoch in range(checkpoint_epochs + 1, args.max_epochs):
        train_loss = 0.
        time_ = time.time()
        model.train()
        for batch_idx, (input,_) in enumerate(tqdm.tqdm(train_loader, desc=f"Train epoch {epoch}")):
            input = input.cuda(non_blocking=True)

            obs = input.shape[1:]
            all_masks = all_masks_by_obs[obs]
            order_i = np.random.randint(len(all_masks))
            mask_init, mask_undilated, mask_dilated = all_masks[order_i]
            output = model(input, mask_init=mask_init, mask_undilated=mask_undilated, mask_dilated=mask_dilated)

            loss = loss_op(input, output)
            deno = args.batch_size * np.prod(obs) * np.log(2.)
            assert deno > 0, embed()
            train_bpd = loss / deno
            if args.minimize_bpd:
                loss = train_bpd

            if batch_idx % args.accum_freq == 0:
                optimizer.zero_grad()
            loss.backward()
            if (batch_idx + 1) % args.accum_freq == 0:
                if args.clip > 0:
                    # Compute and rescale gradient norm
                    gradient_norm = nn.utils.clip_grad_norm_(model.parameters(), args.clip)
                    # if gradient_norm > args.clip:
                        # logger.warning(f"Clipped gradients to norm {args.clip}")
                else:
                    # Just compute the gradient norm
                    parameters = list(filter(lambda p: p.grad is not None, model.parameters()))
                    gradient_norm = 0
                    for p in parameters:
                        param_norm = p.grad.data.norm(2)
                        gradient_norm += param_norm.item() ** 2
                    gradient_norm = gradient_norm ** (1. / 2)
                writer.add_scalar('train/gradient_norm', gradient_norm, global_step)
                optimizer.step()
            train_loss += loss.item()

            writer.add_scalar('train/bpd', train_bpd.item(), global_step)
            min_train_bpd = min(min_train_bpd, train_bpd.item())
            writer.add_scalar('train/min_bpd', min_train_bpd, global_step)

            if batch_idx >= 100 and train_bpd.item() >= 10:
                logger.warning("WARNING: main.py: large batch loss {} bpd".format(train_bpd.item()))

            if (batch_idx + 1) % args.print_every == 0: 
                deno = args.print_every * args.batch_size * np.prod(obs) * np.log(2.)
                average_bpd = train_loss / args.print_every if args.minimize_bpd else train_loss / deno
                logger.info('train bpd : {:.4f}, train loss : {:.1f}, time : {:.4f}, global step: {}'.format(
                    average_bpd,
                    train_loss / args.print_every,
                    (time.time() - time_),
                    global_step))
                train_loss = 0.
                time_ = time.time()

            if (batch_idx + 1) % args.accum_freq == 0:
                global_step += 1

        # decrease learning rate
        scheduler.step()

        model.eval()
        with torch.no_grad():
            save_dict = {}

            for obs in resized_obses:
                logger.info(f"testing with obs {obs2str(obs)}...")
                test_bpd = test(model,
                                all_masks_by_obs[obs],
                                test_loader_by_obs[obs],
                                epoch,
                                progress_bar=True)
                writer.add_scalar(f'test/bpd_{obs2str(obs)}', test_bpd, global_step)
                logger.info(f"test loss for obs {obs2str(obs)}: %s bpd" % test_bpd)
                save_dict[f"test_loss_{obs2str(obs)}"] = test_bpd

                # Log min test bpd for smoothness
                min_test_bpd_by_obs[obs] = min(min_test_bpd_by_obs[obs], test_bpd)
                writer.add_scalar(f'test/min_bpd_{obs2str(obs)}', min_test_bpd_by_obs[obs], global_step)
                if obs == dataset_obs:
                    writer.add_scalar(f'test/bpd', test_bpd, global_step)
                    writer.add_scalar(f'test/min_bpd', min_test_bpd_by_obs[obs], global_step)

            if (epoch + 1) % args.save_interval == 0: 
                logger.info('saving model...')
                save_dict["epoch"] = epoch
                save_dict["args"] = vars(args)
                try:
                    save_dict["model_state_dict"] = model.module.state_dict()
                    save_dict["optimizer_state_dict"] = optimizer.state_dict()
                    torch.save(save_dict, os.path.join(run_dir, f"{args.exp_id}_ep{epoch}.pth"))
                except Exception as e:
                    logger.error("Failed to save checkpoint! Error: %s", e)

            if (epoch + 1) % args.sample_interval == 0: 
                for obs in resized_obses:
                    try:
                        all_masks = all_masks_by_obs[obs]
                        all_generation_idx = all_generation_idx_by_obs[obs]
                        sample_order_i = np.random.randint(len(all_masks))
                        logger.info('sampling images with observation %s, ordering variant %d...', obs2str(obs), sample_order_i)
                        sample_t = sample(model, all_generation_idx[sample_order_i], *all_masks[sample_order_i])
                        sample_t = rescaling_inv(sample_t)
                        utils.save_image(sample_t, os.path.join(run_dir, f"tsample_obs{obs2str(obs)}_{epoch}_order{sample_order_i}.png"), 
                                         nrow=5, padding=0)
                    except Exception as e:
                        logger.error("Failed to sample images! Error: %s", e)
elif args.mode == "sample":
    model.eval()
    with torch.no_grad():
        for obs in resized_obses:
            all_masks = all_masks_by_obs[obs]
            all_generation_idx = all_generation_idx_by_obs[obs]
            sample_order_i = np.random.randint(len(all_masks))
            logger.info('sampling images with observation %s, ordering variant %d...', obs2str(obs), sample_order_i)
            sample_t = sample(model, all_generation_idx[sample_order_i], *all_masks[sample_order_i])
            sample_t = rescaling_inv(sample_t)
            utils.save_image(sample_t, os.path.join(run_dir, f'sample_obs{obs2str(obs)}_{checkpoint_epochs}_order{sample_order_i}.png'),
                             nrow=5, padding=0)
elif args.mode == "test":
    model.eval()
    with torch.no_grad():
        for obs in resized_obses:
            logger.info(f"testing with obs {obs2str(obs)}...")
            test_bpd = test(model,
                            all_masks_by_obs[obs],
                            test_loader_by_obs[obs],
                            checkpoint_epochs,
                            progress_bar=False)
            logger.info(f"test loss for obs {obs2str(obs)}: %s bpd" % test_bpd)