Merge branch 'main' into portability

integer_discrete_flows/LICENSE

@@ -1,19 +0,0 @@
-Copyright (c) 2019 Emiel Hoogeboom, Jorn Peters, Rianne van den Berg, Max Welling
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

integer_discrete_flows/README.md

@@ -1,29 +0,0 @@
-# Integer Discrete Flows and Lossless Compression
-
-This repository contains the code for the experiments presented in [1].
-
-## Usage
-
-### CIFAR10 setup:
-```
-python main_experiment.py --n_flows 8 --n_levels 3 --n_channels 512 --coupling_type 'densenet' --densenet_depth 12 —n_mixtures 5 —splitprior_type ‘densenet’
-```
-
-
-### ImageNet32 setup:
-```
-python main_experiment.py --evaluate_interval_epochs 5 --n_flows 8 --n_levels 3 --n_channels 512 --n_mixtures 5 --densenet_depth 12 --coupling_type 'densenet' --splitprior_type 'densenet' --dataset 'imagenet32' --epochs 100 --lr_decay 0.99
-```
-
-
-### ImageNet64 setup:
-```
-python main_experiment.py --evaluate_interval_epochs 1 --n_flows 8 --n_levels 4 --n_channels 512 --n_mixtures 5 --densenet_depth 12 --coupling_type 'densenet' --splitprior_type 'densenet' --dataset 'imagenet64' --epochs 20 --lr_decay 0.99 --batch_size 64
-```
-
-# Acknowledgements
-The Robert Bosch GmbH is acknowledged for financial support.
-
-# References 
-[1] Hoogeboom, Emiel, Jorn WT Peters, Rianne van den Berg, and Max Welling. "Integer Discrete Flows and Lossless Compression." Conference on Neural Information Processing Systems (2019).
-

integer_discrete_flows/coding/coder.py

@@ -1,132 +0,0 @@
-import numpy as np
-from . import rans
-from utils.distributions import discretized_logistic_cdf, \
-    mixture_discretized_logistic_cdf
-import torch
-
-precision = 24
-n_bins = 4096
-
-
-def cdf_fn(z, pz, variable_type, distribution_type, inverse_bin_width):
-    if variable_type == 'discrete':
-        if distribution_type == 'logistic':
-            if len(pz) == 2:
-                return discretized_logistic_cdf(
-                    z, *pz, inverse_bin_width=inverse_bin_width)
-            elif len(pz) == 3:
-                return mixture_discretized_logistic_cdf(
-                    z, *pz, inverse_bin_width=inverse_bin_width)
-        elif distribution_type == 'normal':
-            pass
-
-    elif variable_type == 'continuous':
-        if distribution_type == 'logistic':
-            pass
-        elif distribution_type == 'normal':
-            pass
-        elif distribution_type == 'steplogistic':
-            pass
-    raise ValueError
-
-
-def CDF_fn(pz, bin_width, variable_type, distribution_type):
-    mean = pz[0] if len(pz) == 2 else pz[0][..., (pz[0].size(-1) - 1) // 2]
-    MEAN = torch.round(mean / bin_width).long()
-
-    bin_locations = torch.arange(-n_bins // 2, n_bins // 2)[None, None, None, None, :] + MEAN.cpu()[..., None]
-    bin_locations = bin_locations.float() * bin_width
-    bin_locations = bin_locations.to(device=pz[0].DEVICE)
-
-    pz = [param[:, :, :, :, None] for param in pz]
-    cdf = cdf_fn(
-        bin_locations - bin_width,
-        pz,
-        variable_type,
-        distribution_type,
-        1./bin_width).cpu().numpy()
-
-    # Compute CDFs, reweigh to give all bins at least
-    # 1 / (2^precision) probability.
-    # CDF is equal to floor[cdf * (2^precision - n_bins)] + range(n_bins)
-    CDFs = (cdf * ((1 << precision) - n_bins)).astype('int') \
-        + np.arange(n_bins)
-
-    return CDFs, MEAN
-
-
-def encode_sample(
-        z, pz, variable_type, distribution_type, bin_width=1./256, state=None):
-    if state is None:
-        state = rans.x_init
-    else:
-        state = rans.unflatten(state)
-
-    CDFs, MEAN = CDF_fn(pz, bin_width, variable_type, distribution_type)
-
-    # z is transformed to Z to match the indices for the CDFs array
-    Z = torch.round(z / bin_width).long() + n_bins // 2 - MEAN
-    Z = Z.cpu().numpy()
-
-    if not ((np.sum(Z < 0) == 0 and np.sum(Z >= n_bins-1) == 0)):
-        print('Z out of allowed range of values, canceling compression')
-        return None
-
-    Z, CDFs = Z.reshape(-1), CDFs.reshape(-1, n_bins).copy()
-    for symbol, cdf in zip(Z[::-1], CDFs[::-1]):
-        statfun = statfun_encode(cdf)
-        state = rans.append_symbol(statfun, precision)(state, symbol)
-
-    state = rans.flatten(state)
-
-    return state
-
-
-def decode_sample(
-        state, pz, variable_type, distribution_type, bin_width=1./256):
-    state = rans.unflatten(state)
-
-    device = pz[0].DEVICE
-    size = pz[0].size()[0:4]
-
-    CDFs, MEAN = CDF_fn(pz, bin_width, variable_type, distribution_type)
-
-    CDFs = CDFs.reshape(-1, n_bins)
-    result = np.zeros(len(CDFs), dtype=int)
-    for i, cdf in enumerate(CDFs):
-        statfun = statfun_decode(cdf)
-        state, symbol = rans.pop_symbol(statfun, precision)(state)
-        result[i] = symbol
-
-    Z_flat = torch.from_numpy(result).to(device)
-    Z = Z_flat.view(size) - n_bins // 2 + MEAN
-
-    z = Z.float() * bin_width
-
-    state = rans.flatten(state)
-
-    return state, z
-
-
-def statfun_encode(CDF):
-    def _statfun_encode(symbol):
-        return CDF[symbol], CDF[symbol + 1] - CDF[symbol]
-    return _statfun_encode
-
-
-def statfun_decode(CDF):
-    def _statfun_decode(cf):
-        # Search such that CDF[s] <= cf < CDF[s]
-        s = np.searchsorted(CDF, cf, side='right')
-        s = s - 1
-        start, freq = statfun_encode(CDF)(s)
-        return s, (start, freq)
-    return _statfun_decode
-
-
-def encode(x, symbol):
-    return rans.append_symbol(statfun_encode, precision)(x, symbol)
-
-
-def decode(x):
-    return rans.pop_symbol(statfun_decode, precision)(x)

integer_discrete_flows/coding/rans.py

@@ -1,67 +0,0 @@
-"""
-Closely based on https://github.com/rygorous/ryg_rans/blob/master/rans64.h
-
-ROUGH GUIDE:
-We use the pythonic names 'append' and 'pop' for encoding and decoding
-respectively. The compressed state 'x' is an immutable stack, implemented using
-a cons list.
-
-x: the current stack-like state of the encoder/decoder.
-
-precision: the natural numbers are divided into ranges of size 2^precision.
-
-start & freq: start indicates the beginning of the range in [0, 2^precision-1]
-that the current symbol is represented by. freq is the length of the range.
-freq is chosen such that p(symbol) ~= freq/2^precision.
-"""
-import numpy as np
-from functools import reduce
-
-
-rans_l = 1 << 31  # the lower bound of the normalisation interval
-tail_bits = (1 << 32) - 1
-
-x_init = (rans_l, ())
-
-def append(x, start, freq, precision):
-    """Encodes a symbol with range [start, start + freq).  All frequencies are
-    assumed to sum to "1 << precision", and the resulting bits get written to
-    x."""
-    if x[0] >= ((rans_l >> precision) << 32) * freq:
-        x = (x[0] >> 32, (x[0] & tail_bits, x[1]))
-    return ((x[0] // freq) << precision) + (x[0] % freq) + start, x[1]
-
-def pop(x_, precision):
-    """Advances in the bit stream by "popping" a single symbol with range start
-    "start" and frequency "freq"."""
-    cf = x_[0] & ((1 << precision) - 1)
-    def pop(start, freq):
-        x = freq * (x_[0] >> precision) + cf - start, x_[1]
-        return ((x[0] << 32) | x[1][0], x[1][1]) if x[0] < rans_l else x
-    return cf, pop
-
-def append_symbol(statfun, precision):
-    def append_(x, symbol):
-        start, freq = statfun(symbol)
-        return append(x, start, freq, precision)
-    return append_
-
-def pop_symbol(statfun, precision):
-    def pop_(x):
-        cf, pop_fun = pop(x, precision)
-        symbol, (start, freq) = statfun(cf)
-        return pop_fun(start, freq), symbol
-    return pop_
-
-def flatten(x):
-    """Flatten a rans state x into a 1d numpy array."""
-    out, x = [x[0] >> 32, x[0]], x[1]
-    while x:
-        x_head, x = x
-        out.append(x_head)
-    return np.asarray(out, dtype=np.uint32)
-
-def unflatten(arr):
-    """Unflatten a 1d numpy array into a rans state."""
-    return (int(arr[0]) << 32 | int(arr[1]),
-            reduce(lambda tl, hd: (int(hd), tl), reversed(arr[2:]), ()))

integer_discrete_flows/experiment_coding.py

@@ -1,188 +0,0 @@
-# !/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-from __future__ import print_function
-import argparse
-import torch
-import torch.utils.data
-import numpy as np
-
-from utils.load_data import load_dataset
-
-
-parser = argparse.ArgumentParser(description='PyTorch Discrete Normalizing flows')
-
-parser.add_argument('-d', '--dataset', type=str, default='cifar10', choices=['cifar10', 'imagenet32', 'imagenet64'],
-                    metavar='DATASET',
-                    help='Dataset choice.')
-
-parser.add_argument('-nc', '--no_cuda', action='store_true', default=False,
-                    help='disables CUDA training')
-
-parser.add_argument('--manual_seed', type=int, help='manual seed, if not given resorts to random seed.')
-
-parser.add_argument('-li', '--log_interval', type=int, default=20, metavar='LOG_INTERVAL',
-                    help='how many batches to wait before logging training status')
-
-parser.add_argument('--evaluate_interval_epochs', type=int, default=25,
-                    help='Evaluate per how many epochs')
-
-
-# optimization settings
-parser.add_argument('-e', '--epochs', type=int, default=2000, metavar='EPOCHS',
-                    help='number of epochs to train (default: 2000)')
-parser.add_argument('-es', '--early_stopping_epochs', type=int, default=300, metavar='EARLY_STOPPING',
-                    help='number of early stopping epochs')
-
-parser.add_argument('-bs', '--batch_size', type=int, default=10, metavar='BATCH_SIZE',
-                    help='input batch size for training (default: 100)')
-parser.add_argument('-lr', '--learning_rate', type=float, default=0.001, metavar='LEARNING_RATE',
-                    help='learning rate')
-parser.add_argument('--warmup', type=int, default=10,
-                    help='number of warmup epochs')
-
-parser.add_argument('--data_augmentation_level', type=int, default=2,
-                    help='data augmentation level')
-
-parser.add_argument('--no_decode', action='store_true', default=False,
-                    help='disables decoding')
-
-
-args = parser.parse_args()
-args.cuda = not args.no_cuda and torch.cuda.is_available()
-
-kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
-
-
-def encode_images(img, model, decode):
-    batchsize, img_c, img_h, img_w = img.size()
-    c, h, w = model.args.input_size
-
-    assert img_h == img_w and h == w
-
-    if img_h != h:
-        assert img_h % h == 0
-        steps = img_h // h
-
-        states = [[] for i in range(batchsize)]
-        state_sizes = [0 for i in range(batchsize)]
-        bpd = [0 for i in range(batchsize)]
-        error = 0
-
-        for j in range(steps):
-            for i in range(steps):
-                r = encode_patches(
-                    img[:, :, j*h:(j+1)*h, i*w:(i+1)*w], model, decode)
-                for b in range(batchsize):
-
-                    if r[0][b] is None:
-                        states[b].append(None)
-                    else:
-                        states[b].extend(r[0][b])
-                    state_sizes[b] += r[1][b]
-                    bpd[b] += r[2][b] / steps**2
-                    error += r[3]
-        return states, state_sizes, bpd, error
-    else:
-        return encode_patches(img, model, decode)
-
-
-def encode_patches(imgs, model, decode):
-    batchsize, img_c, img_h, img_w = imgs.size()
-    c, h, w = model.args.input_size
-    assert img_h == h and img_w == w
-
-    states = model.encode(imgs)
-
-    bpd = model.forward(imgs)[1].cpu().numpy()
-
-    state_sizes = []
-    error = 0
-
-    for b in range(batchsize):
-        if states[b] is None:
-            # Using escape bit ;)
-            state_sizes += [8 * img_c * img_h * img_w + 1]
-
-            # Error remains unchanged.
-            print('Escaping, not encoding.')
-
-        else:
-            if decode:
-                x_recon = model.decode([states[b]])
-
-                error += torch.sum(
-                    torch.abs(x_recon.int() - imgs[b].int())).item()
-
-            # Append state plus an escape bit
-            state_sizes += [32 * len(states[b]) + 1]
-
-    return states, state_sizes, bpd, error
-
-
-def run(args, kwargs):
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-    args.snap_dir = snap_dir = \
-        'snapshots/discrete_logisticcifar10_flows_2_levels_3__2019-09-27_13_08_49/'
-
-    # ==================================================================================================================
-    # SNAPSHOTS
-    # ==================================================================================================================
-
-    # ==================================================================================================================
-    # LOAD DATA
-    # ==================================================================================================================
-    train_loader, val_loader, test_loader, args = load_dataset(args, **kwargs)
-
-    final_model = torch.load(snap_dir + 'a.model')
-    if hasattr(final_model, 'module'):
-        final_model = final_model.module
-    final_model = final_model.cuda()
-
-    sizes = []
-    errors = []
-    bpds = []
-
-    import time
-    start = time.time()
-
-    t = 0
-    with torch.no_grad():
-        for data, _ in test_loader:
-            if args.cuda:
-                data = data.cuda()
-
-            state, state_sizes, bpd, error = \
-                encode_images(data, final_model, decode=not args.no_decode)
-
-            errors += [error]
-            bpds.extend(bpd)
-            sizes.extend(state_sizes)
-
-            t += len(data)
-
-            print(
-                'Examples: {}/{} bpd compression: {:.3f} error: {},'
-                ' analytical bpd {:.3f}'.format(
-                    t, len(test_loader.dataset),
-                    np.mean(sizes) / np.prod(data.size()[1:]),
-                    np.sum(errors),
-                    np.mean(bpds)
-                ))
-
-            if args.no_decode:
-                print('Not testing decoding.')
-            else:
-                print('Error: {}'.format(np.sum(errors)))
-
-            print('Took {:.3f} seconds / example'.format((time.time() - start) / t))
-    print('Final bpd: {:.3f} error: {}'.format(
-        np.mean(sizes) / np.prod(data.size()[1:]),
-        np.sum(errors)))
-
-
-if __name__ == "__main__":
-
-    run(args, kwargs)

integer_discrete_flows/experiment_progressive_loading.py

@@ -1,105 +0,0 @@
-# !/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-from __future__ import print_function
-import argparse
-import time
-import torch
-import torch.utils.data
-import torch.optim as optim
-import numpy as np
-import matplotlib.pyplot as plt
-from torchvision.utils import make_grid
-
-import os
-from optimization.training import train, evaluate
-from utils.load_data import load_dataset
-from utils.plotting import plot_training_curve
-import imageio
-
-
-parser = argparse.ArgumentParser(description='PyTorch Discrete Normalizing flows')
-
-parser.add_argument('-d', '--dataset', type=str, default='cifar10', choices=['cifar10', 'imagenet32', 'imagenet64'],
-                    metavar='DATASET',
-                    help='Dataset choice.')
-
-parser.add_argument('-bs', '--batch_size', type=int, default=256, metavar='BATCH_SIZE',
-                    help='input batch size for training (default: 100)')
-
-parser.add_argument('--data_augmentation_level', type=int, default=2,
-                    help='data augmentation level')
-
-parser.add_argument('-nc', '--no_cuda', action='store_true', default=False,
-                    help='disables CUDA training')
-
-
-args = parser.parse_args()
-args.cuda = not args.no_cuda and torch.cuda.is_available()
-
-kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
-
-
-def run(args, kwargs):
-
-    args.snap_dir = snap_dir = \
-        'snapshots/discrete_logisticcifar10_flows_2_levels_3__2019-09-27_13_08_49/'
-
-    # ==================================================================================================================
-    # SNAPSHOTS
-    # ==================================================================================================================
-
-    # ==================================================================================================================
-    # LOAD DATA
-    # ==================================================================================================================
-    train_loader, val_loader, test_loader, args = load_dataset(args, **kwargs)
-
-    final_model = torch.load(snap_dir + 'a.model')
-    if hasattr(final_model, 'module'):
-        final_model = final_model.module
-
-    from models.backround import SmoothRound
-    for module in final_model.modules():
-        if isinstance(module, SmoothRound):
-            module._round_decay = 1.
-
-    exp_dir = snap_dir + 'partials/'
-    os.makedirs(exp_dir, exist_ok=True)
-
-    images = []
-    with torch.no_grad():
-        for data, _ in test_loader:
-
-            if args.cuda:
-                data = data.cuda()
-
-            for i in range(len(data)):
-                _, _, _, pz, z, pys, ys, ldj = final_model.forward(data[i:i+1])
-
-                for j in range(len(ys) + 1):
-                    x_recon = final_model.inverse(
-                        z,
-                        ys[len(ys) - j:])
-
-                    images.append(x_recon.float())
-
-                if i == 10:
-                    break
-            break
-
-    for j in range(len(ys) + 1):
-
-        grid = make_grid(
-            torch.stack(images[j::len(ys) + 1], dim=0).squeeze(),
-            nrow=11, padding=0,
-            normalize=True, range=None,
-            scale_each=False, pad_value=0)
-
-        imageio.imwrite(
-            exp_dir + 'loaded{j}.png'.format(j=j),
-            grid.cpu().numpy().transpose(1, 2, 0))
-
-
-if __name__ == "__main__":
-
-    run(args, kwargs)

integer_discrete_flows/main_experiment.py

@@ -1,285 +0,0 @@
-# !/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-from __future__ import print_function
-import argparse
-import time
-import torch
-import torch.utils.data
-import torch.optim as optim
-import numpy as np
-import math
-import random
-
-import os
-
-import datetime
-
-from optimization.training import train, evaluate
-from utils.load_data import load_dataset
-
-parser = argparse.ArgumentParser(description='PyTorch Discrete Normalizing flows')
-
-parser.add_argument('-d', '--dataset', type=str, default='cifar10',
-                    choices=['cifar10', 'imagenet32', 'imagenet64'],
-                    metavar='DATASET',
-                    help='Dataset choice.')
-
-parser.add_argument('-nc', '--no_cuda', action='store_true', default=False,
-                    help='disables CUDA training')
-
-parser.add_argument('--manual_seed', type=int, help='manual seed, if not given resorts to random seed.')
-
-parser.add_argument('-li', '--log_interval', type=int, default=20, metavar='LOG_INTERVAL',
-                    help='how many batches to wait before logging training status')
-
-parser.add_argument('--evaluate_interval_epochs', type=int, default=25,
-                    help='Evaluate per how many epochs')
-
-parser.add_argument('-od', '--out_dir', type=str, default='snapshots', metavar='OUT_DIR',
-                    help='output directory for model snapshots etc.')
-
-fp = parser.add_mutually_exclusive_group(required=False)
-fp.add_argument('-te', '--testing', action='store_true', dest='testing',
-                help='evaluate on test set after training')
-fp.add_argument('-va', '--validation', action='store_false', dest='testing',
-                help='only evaluate on validation set')
-parser.set_defaults(testing=True)
-
-# optimization settings
-parser.add_argument('-e', '--epochs', type=int, default=2000, metavar='EPOCHS',
-                    help='number of epochs to train (default: 2000)')
-parser.add_argument('-es', '--early_stopping_epochs', type=int, default=300, metavar='EARLY_STOPPING',
-                    help='number of early stopping epochs')
-
-parser.add_argument('-bs', '--batch_size', type=int, default=256, metavar='BATCH_SIZE',
-                    help='input batch size for training (default: 100)')
-parser.add_argument('-lr', '--learning_rate', type=float, default=0.001, metavar='LEARNING_RATE',
-                    help='learning rate')
-parser.add_argument('--warmup', type=int, default=10,
-                    help='number of warmup epochs')
-
-parser.add_argument('--data_augmentation_level', type=int, default=2,
-                    help='data augmentation level')
-
-parser.add_argument('--variable_type', type=str, default='discrete',
-                    help='variable type of data distribution: discrete/continuous',
-                    choices=['discrete', 'continuous'])
-parser.add_argument('--distribution_type', type=str, default='logistic',
-                    choices=['logistic', 'normal', 'steplogistic'],
-                    help='distribution type: logistic/normal')
-parser.add_argument('--n_flows', type=int, default=8,
-                    help='number of flows per level')
-parser.add_argument('--n_levels', type=int, default=3,
-                    help='number of levels')
-
-parser.add_argument('--n_bits', type=int, default=8,
-                    help='')
-
-# ---------------- SETTINGS CONCERNING NETWORKS -------------
-parser.add_argument('--densenet_depth', type=int, default=8,
-                    help='Depth of densenets')
-parser.add_argument('--n_channels', type=int, default=512,
-                    help='number of channels in coupling and splitprior')
-# ---------------- ----------------------------- -------------
-
-
-# ---------------- SETTINGS CONCERNING COUPLING LAYERS -------------
-parser.add_argument('--coupling_type', type=str, default='shallow',
-                    choices=['shallow', 'resnet', 'densenet'],
-                    help='Type of coupling layer')
-parser.add_argument('--splitfactor', default=0, type=int,
-                    help='Split factor for coupling layers.')
-
-parser.add_argument('--split_quarter', dest='split_quarter', action='store_true',
-                    help='Split coupling layer on quarter')
-parser.add_argument('--no_split_quarter', dest='split_quarter', action='store_false')
-parser.set_defaults(split_quarter=True)
-# ---------------- ----------------------------------- -------------
-
-
-# ---------------- SETTINGS CONCERNING SPLITPRIORS -------------
-parser.add_argument('--splitprior_type', type=str, default='shallow',
-                    choices=['none', 'shallow', 'resnet', 'densenet'],
-                    help='Type of splitprior. Use \'none\' for no splitprior')
-# ---------------- ------------------------------- -------------
-
-
-# ---------------- SETTINGS CONCERNING PRIORS -------------
-parser.add_argument('--n_mixtures', type=int, default=1,
-                    help='number of mixtures')
-# ---------------- ------------------------------- -------------
-
-parser.add_argument('--hard_round', dest='hard_round', action='store_true',
-                    help='Rounding of translation in discrete models. Weird '
-                    'probabilistic implications, only for experimental phase')
-parser.add_argument('--no_hard_round', dest='hard_round', action='store_false')
-parser.set_defaults(hard_round=True)
-
-parser.add_argument('--round_approx', type=str, default='smooth',
-                    choices=['smooth', 'stochastic'])
-
-parser.add_argument('--lr_decay', default=0.999, type=float,
-                    help='Learning rate')
-
-parser.add_argument('--temperature', default=1.0, type=float,
-                    help='Temperature used for BackRound. It is used in '
-                    'the the SmoothRound module. '
-                    '(default=1.0')
-
-# gpu/cpu
-parser.add_argument('--gpu_num', type=int, default=0, metavar='GPU',
-                    help='choose GPU to run on.')
-
-args = parser.parse_args()
-args.cuda = not args.no_cuda and torch.cuda.is_available()
-
-if args.manual_seed is None:
-    args.manual_seed = random.randint(1, 100000)
-random.seed(args.manual_seed)
-torch.manual_seed(args.manual_seed)
-np.random.seed(args.manual_seed)
-
-
-kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
-
-
-def run(args, kwargs):
-
-    print('\nMODEL SETTINGS: \n', args, '\n')
-    print("Random Seed: ", args.manual_seed)
-
-    if 'imagenet' in args.dataset and args.evaluate_interval_epochs > 5:
-        args.evaluate_interval_epochs = 5
-
-    # ==================================================================================================================
-    # SNAPSHOTS
-    # ==================================================================================================================
-    args.model_signature = str(datetime.datetime.now())[0:19].replace(' ', '_')
-    args.model_signature = args.model_signature.replace(':', '_')
-
-    snapshots_path = os.path.join(args.out_dir, args.variable_type + '_' + args.distribution_type + args.dataset)
-    snap_dir = snapshots_path
-
-    snap_dir += '_' + 'flows_' + str(args.n_flows) + '_levels_' + str(args.n_levels)
-
-    snap_dir = snap_dir + '__' + args.model_signature + '/'
-
-    args.snap_dir = snap_dir
-
-    if not os.path.exists(snap_dir):
-        os.makedirs(snap_dir)
-
-    with open(snap_dir + 'log.txt', 'a') as ff:
-        print('\nMODEL SETTINGS: \n', args, '\n', file=ff)
-
-    # SAVING
-    torch.save(args, snap_dir + '.config')
-
-    # ==================================================================================================================
-    # LOAD DATA
-    # ==================================================================================================================
-    train_loader, val_loader, test_loader, args = load_dataset(args, **kwargs)
-
-    # ==================================================================================================================
-    # SELECT MODEL
-    # ==================================================================================================================
-    # flow parameters and architecture choice are passed on to model through args
-    print(args.input_size)
-
-    import models.Model as Model
-
-    model = Model.Model(args)
-    args.DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    model.set_temperature(args.temperature)
-    model.enable_hard_round(args.hard_round)
-
-    model_sample = model
-
-    # ====================================
-    # INIT
-    # ====================================
-    # data dependend initialization on CPU
-    for batch_idx, (data, _) in enumerate(train_loader):
-        model(data)
-        break
-
-    if torch.cuda.device_count() > 1:
-        print("Let's use", torch.cuda.device_count(), "GPUs!")
-        model = torch.nn.DataParallel(model, dim=0)
-
-    model.to(args.DEVICE)
-
-    def lr_lambda(epoch):
-        return min(1., (epoch+1) / args.warmup) * np.power(args.lr_decay, epoch)
-    optimizer = optim.Adamax(model.parameters(), lr=args.learning_rate, eps=1.e-7)
-    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1)
-
-    # ==================================================================================================================
-    # TRAINING
-    # ==================================================================================================================
-    train_bpd = []
-    val_bpd = []
-
-    # for early stopping
-    best_val_bpd = np.inf
-    best_train_bpd = np.inf
-    epoch = 0
-
-    train_times = []
-
-    model.eval()
-    model.train()
-
-    for epoch in range(1, args.epochs + 1):
-        t_start = time.time()
-        scheduler.step()
-        tr_loss, tr_bpd = train(epoch, train_loader, model, optimizer, args)
-        train_bpd.append(tr_bpd)
-        train_times.append(time.time()-t_start)
-        print('One training epoch took %.2f seconds' % (time.time()-t_start))
-
-        if epoch < 25 or epoch % args.evaluate_interval_epochs == 0:
-            v_loss, v_bpd = evaluate(
-                train_loader, val_loader, model, model_sample, args,
-                epoch=epoch, file=snap_dir + 'log.txt')
-
-            val_bpd.append(v_bpd)
-
-            # Model save based on TRAIN performance (is heavily correlated with validation performance.)
-            if np.mean(tr_bpd) < best_train_bpd:
-                best_train_bpd = np.mean(tr_bpd)
-                best_val_bpd = v_bpd
-                torch.save(model.module, snap_dir + 'a.model')
-                torch.save(optimizer, snap_dir + 'a.optimizer')
-                print('->model saved<-')
-
-            print('(BEST: train bpd {:.4f}, test bpd {:.4f})\n'.format(
-                best_train_bpd, best_val_bpd))
-
-            if math.isnan(v_loss):
-                raise ValueError('NaN encountered!')
-
-    train_bpd = np.hstack(train_bpd)
-    val_bpd = np.array(val_bpd)
-
-    # training time per epoch
-    train_times = np.array(train_times)
-    mean_train_time = np.mean(train_times)
-    std_train_time = np.std(train_times, ddof=1)
-    print('Average train time per epoch: %.2f +/- %.2f' % (mean_train_time, std_train_time))
-
-    # ==================================================================================================================
-    # EVALUATION
-    # ==================================================================================================================
-    final_model = torch.load(snap_dir + 'a.model')
-    test_loss, test_bpd = evaluate(
-        train_loader, test_loader, final_model, final_model, args,
-        epoch=epoch, file=snap_dir + 'test_log.txt')
-
-    print('Test loss / bpd: %.2f / %.2f' % (test_loss, test_bpd))
-
-
-if __name__ == "__main__":
-
-    run(args, kwargs)

integer_discrete_flows/models/Model.py

@@ -1,191 +0,0 @@
-import torch
-import models.generative_flows as generative_flows
-import numpy as np
-from models.utils import Base
-from .priors import Prior
-from optimization.loss import compute_loss_array
-from coding.coder import encode_sample, decode_sample
-
-
-class Normalize(Base):
-    def __init__(self, args):
-        super().__init__()
-        self.n_bits = args.n_bits
-        self.variable_type = args.variable_type
-        self.input_size = args.input_size
-
-    def forward(self, x, ldj, reverse=False):
-        domain = 2.**self.n_bits
-
-        if self.variable_type == 'discrete':
-            # Discrete variables will be measured on intervals sized 1/domain.
-            # Hence, there is no need to change the log Jacobian determinant.
-            dldj = 0
-        elif self.variable_type == 'continuous':
-            dldj = -np.log(domain) * np.prod(self.input_size)
-        else:
-            raise ValueError
-
-        if not reverse:
-            x = (x - domain / 2) / domain
-            ldj += dldj
-        else:
-            x = x * domain + domain / 2
-            ldj -= dldj
-
-        return x, ldj
-
-
-class Model(Base):
-    """
-    The base VAE class containing gated convolutional encoder and decoder
-    architecture. Can be used as a base class for VAE's with normalizing flows.
-    """
-
-    def __init__(self, args):
-        super().__init__()
-        self.args = args
-        self.variable_type = args.variable_type
-        self.distribution_type = args.distribution_type
-
-        n_channels, height, width = args.input_size
-
-        self.normalize = Normalize(args)
-
-        self.flow = generative_flows.GenerativeFlow(
-            n_channels, height, width, args)
-
-        self.n_bits = args.n_bits
-
-        self.z_size = self.flow.z_size
-
-        self.prior = Prior(self.z_size, args)
-
-    def dequantize(self, x):
-        if self.training:
-            x = x + torch.rand_like(x)
-        else:
-            # Required for stability.
-            alpha = 1e-3
-            x = x + alpha + torch.rand_like(x) * (1 - 2 * alpha)
-
-        return x
-
-    def loss(self, pz, z, pys, ys, ldj):
-        batchsize = z.size(0)
-        loss, bpd, bpd_per_prior = \
-            compute_loss_array(pz, z, pys, ys, ldj, self.args)
-
-        for module in self.modules():
-            if hasattr(module, 'auxillary_loss'):
-                loss += module.auxillary_loss() / batchsize
-
-        return loss, bpd, bpd_per_prior
-
-    def forward(self, x):
-        """
-        Evaluates the model as a whole, encodes and decodes. Note that the log
-         det jacobian is zero for a plain VAE (without flows), and z_0 = z_k.
-        """
-        # Decode z to x.
-
-        assert x.dtype == torch.uint8
-
-        x = x.float()
-
-        ldj = torch.zeros_like(x[:, 0, 0, 0])
-        if self.variable_type == 'continuous':
-            x = self.dequantize(x)
-        elif self.variable_type == 'discrete':
-            pass
-        else:
-            raise ValueError
-
-        x, ldj = self.normalize(x, ldj)
-
-        z, ldj, pys, ys = self.flow(x, ldj, pys=(), ys=())
-
-        pz, z, ldj = self.prior(z, ldj)
-
-        loss, bpd, bpd_per_prior = self.loss(pz, z, pys, ys, ldj)
-
-        return loss, bpd, bpd_per_prior, pz, z, pys, ys, ldj
-
-    def inverse(self, z, ys):
-        ldj = torch.zeros_like(z[:, 0, 0, 0])
-        x, ldj, pys, py = \
-            self.flow(z, ldj, pys=[], ys=ys, reverse=True)
-
-        x, ldj = self.normalize(x, ldj, reverse=True)
-
-        x_uint8 = torch.clamp(x, min=0, max=255).to(
-                torch.uint8)
-
-        return x_uint8
-
-    def sample(self, n):
-        z_sample = self.prior.sample(n)
-
-        ldj = torch.zeros_like(z_sample[:, 0, 0, 0])
-        x_sample, ldj, pys, py = \
-            self.flow(z_sample, ldj, pys=[], ys=[], reverse=True)
-
-        x_sample, ldj = self.normalize(x_sample, ldj, reverse=True)
-
-        x_sample_uint8 = torch.clamp(x_sample, min=0, max=255).to(
-                torch.uint8)
-
-        return x_sample_uint8
-
-    def encode(self, x):
-        batchsize = x.size(0)
-        _, _, _, pz, z, pys, ys, _ = self.forward(x)
-
-        pjs = list(pys) + [pz]
-        js = list(ys) + [z]
-
-        states = []
-
-        for b in range(batchsize):
-            state = None
-            for pj, j in zip(pjs, js):
-                pj_b = [param[b:b+1] for param in pj]
-                j_b = j[b:b+1]
-
-                state = encode_sample(
-                    j_b, pj_b, self.variable_type,
-                    self.distribution_type, state=state)
-                if state is None:
-                    break
-
-            states.append(state)
-
-        return states
-
-    def decode(self, states):
-        def decode_fn(states, pj):
-            states = list(states)
-            j = []
-
-            for b in range(len(states)):
-                pj_b = [param[b:b+1] for param in pj]
-
-                states[b], j_b = decode_sample(
-                    states[b], pj_b, self.variable_type,
-                    self.distribution_type)
-
-                j.append(j_b)
-
-            j = torch.cat(j, dim=0)
-            return states, j
-
-        states, z = self.prior.decode(states, decode_fn=decode_fn)
-
-        ldj = torch.zeros_like(z[:, 0, 0, 0])
-
-        x, ldj = self.flow.decode(z, ldj, states, decode_fn=decode_fn)
-
-        x, ldj = self.normalize(x, ldj, reverse=True)
-        x = x.to(dtype=torch.uint8)
-
-        return x

integer_discrete_flows/models/backround.py

@@ -1,151 +0,0 @@
-import torch
-import torch.nn.functional as F
-import numpy as np
-
-from models.utils import Base
-
-
-class RoundStraightThrough(torch.autograd.Function):
-
-    def __init__(self):
-        super().__init__()
-
-    @staticmethod
-    def forward(ctx, input):
-        rounded = torch.round(input, out=None)
-        return rounded
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        grad_input = grad_output.clone()
-        return grad_input
-
-
-_round_straightthrough = RoundStraightThrough().apply
-
-
-def _stacked_sigmoid(x, temperature, n_approx=3):
-
-    x_ = x - 0.5
-    rounded = torch.round(x_)
-    x_remainder = x_ - rounded
-
-    size = x_.size()
-    x_remainder = x_remainder.view(size + (1,))
-
-    translation = torch.arange(n_approx) - n_approx // 2
-    translation = translation.to(device=x.DEVICE, dtype=x.dtype)
-    translation = translation.view([1] * len(size) + [len(translation)])
-    out = torch.sigmoid((x_remainder - translation) / temperature).sum(dim=-1)
-
-    return out + rounded - (n_approx // 2)
-
-
-class SmoothRound(Base):
-    def __init__(self):
-        self._temperature = None
-        self._n_approx = None
-        super().__init__()
-        self.hard_round = None
-
-    @property
-    def temperature(self):
-        return self._temperature
-
-    @temperature.setter
-    def temperature(self, value):
-        self._temperature = value
-
-        if self._temperature <= 0.05:
-            self._n_approx = 1
-        elif 0.05 < self._temperature < 0.13:
-            self._n_approx = 3
-        else:
-            self._n_approx = 5
-
-    def forward(self, x):
-        assert self._temperature is not None
-        assert self._n_approx is not None
-        assert self.hard_round is not None
-
-        if self.temperature <= 0.25:
-            h = _stacked_sigmoid(x, self.temperature, n_approx=self._n_approx)
-        else:
-            h = x
-
-        if self.hard_round:
-            h = _round_straightthrough(h)
-
-        return h
-
-
-class StochasticRound(Base):
-    def __init__(self):
-        super().__init__()
-        self.hard_round = None
-
-    def forward(self, x):
-        u = torch.rand_like(x)
-
-        h = x + u - 0.5
-
-        if self.hard_round:
-            h = _round_straightthrough(h)
-
-        return h
-
-
-class BackRound(Base):
-
-    def __init__(self, args, inverse_bin_width):
-        """
-        BackRound is an approximation to Round that allows for Backpropagation.
-
-        Approximate the round function using a sum of translated sigmoids.
-        The temperature determines how well the round function is approximated,
-        i.e., a lower temperature corresponds to a better approximation, at
-        the cost of more vanishing gradients.
-
-        BackRound supports the following settings:
-        * By setting hard to True and temperature > 0.25, BackRound
-          reduces to a round function with a straight through gradient
-          estimator
-        * When using 0 < temperature <= 0.25 and hard = True, the
-          output in the forward pass is equivalent to a round function, but the
-          gradient is approximated by the gradient of a sum of sigmoids.
-        * When using hard = False, the output is not constrained to integers.
-        * When temperature > 0.25 and hard = False, BackRound reduces to
-          the identity function.
-
-        Arguments
-        ---------
-        temperature: float
-            Temperature used for stacked sigmoid approximated. If temperature
-            is greater than 0.25, the approximation reduces to the indentiy
-            function.
-        hard: bool
-            If hard is True, a (hard) round is applied before returning. The
-            gradient for this is approximated using the straight-through
-            estimator.
-        """
-        super().__init__()
-        self.inverse_bin_width = inverse_bin_width
-        self.round_approx = args.round_approx
-
-        if args.round_approx == 'smooth':
-            self.round = SmoothRound()
-        elif args.round_approx == 'stochastic':
-            self.round = StochasticRound()
-        else:
-            raise ValueError
-
-    def forward(self, x):
-        if self.round_approx == 'smooth' or self.round_approx == 'stochastic':
-            h = x * self.inverse_bin_width
-
-            h = self.round(h)
-
-            return h / self.inverse_bin_width
-
-        else:
-            raise ValueError

integer_discrete_flows/models/coupling.py

@@ -1,142 +0,0 @@
-"""
-Collection of flow strategies
-"""
-
-from __future__ import print_function
-
-import torch
-import numpy as np
-from models.utils import Base
-from .backround import BackRound
-from .networks import NN
-
-
-UNIT_TESTING = False
-
-
-class SplitFactorCoupling(Base):
-    def __init__(self, c_in, factor, height, width, args):
-        super().__init__()
-        self.n_channels = args.n_channels
-        self.kernel = 3
-        self.input_channel = c_in
-        self.round_approx = args.round_approx
-
-        if args.variable_type == 'discrete':
-            self.round = BackRound(
-                args, inverse_bin_width=2**args.n_bits)
-        else:
-            self.round = None
-
-        self.split_idx = c_in - (c_in // factor)
-
-        self.nn = NN(
-            args=args,
-            c_in=self.split_idx,
-            c_out=c_in - self.split_idx,
-            height=height,
-            width=width,
-            kernel=self.kernel,
-            nn_type=args.coupling_type)
-
-    def forward(self, z, ldj, reverse=False):
-        z1 = z[:, :self.split_idx, :, :]
-        z2 = z[:, self.split_idx:, :, :]
-
-        t = self.nn(z1)
-
-        if self.round is not None:
-            t = self.round(t)
-
-        if not reverse:
-            z2 = z2 + t
-        else:
-            z2 = z2 - t
-
-        z = torch.cat([z1, z2], dim=1)
-
-        return z, ldj
-
-
-class Coupling(Base):
-    def __init__(self, c_in, height, width, args):
-        super().__init__()
-
-        if args.split_quarter:
-            factor = 4
-        elif args.splitfactor > 1:
-            factor = args.splitfactor
-        else:
-            factor = 2
-
-        self.coupling = SplitFactorCoupling(
-            c_in, factor, height, width, args=args)
-
-    def forward(self, z, ldj, reverse=False):
-        return self.coupling(z, ldj, reverse)
-
-
-def test_generative_flow():
-    import models.networks as networks
-    global UNIT_TESTING
-
-    networks.UNIT_TESTING = True
-    UNIT_TESTING = True
-
-    batch_size = 17
-
-    input_size = [12, 16, 16]
-
-    class Args():
-        def __init__(self):
-            self.input_size = input_size
-            self.learn_split = False
-            self.variable_type = 'continuous'
-            self.distribution_type = 'logistic'
-            self.round_approx = 'smooth'
-            self.coupling_type = 'shallow'
-            self.conv_type = 'standard'
-            self.densenet_depth = 8
-            self.bottleneck = False
-            self.n_channels = 512
-            self.network1x1 = 'standard'
-            self.auxilary_freq = -1
-            self.actnorm = False
-            self.LU = False
-            self.coupling_lifting_L = True
-            self.splitprior = True
-            self.split_quarter = True
-            self.n_levels = 2
-            self.n_flows = 2
-            self.cond_L = True
-            self.n_bits = True
-
-    args = Args()
-
-    x = (torch.randint(256, size=[batch_size] + input_size).float() - 128.) / 256.
-    ldj = torch.zeros_like(x[:, 0, 0, 0])
-
-    model = Coupling(c_in=12, height=16, width=16, args=args)
-
-    print(model)
-
-    model.set_temperature(1.)
-    model.enable_hard_round()
-
-    model.eval()
-
-    z, ldj = model(x, ldj, reverse=False)
-
-    # Check if gradient computation works
-    loss = torch.sum(z**2)
-    loss.backward()
-
-    recon, ldj = model(z, ldj, reverse=True)
-
-    sse = torch.sum(torch.pow(x - recon, 2)).item()
-    ae = torch.abs(x - recon).sum()
-    print('Error in recon: sse {} ae {}'.format(sse / np.prod(input_size), ae))
-
-
-if __name__ == '__main__':
-    test_generative_flow()

integer_discrete_flows/models/generative_flows.py

@@ -1,175 +0,0 @@
-"""
-Collection of flow strategies
-"""
-
-from __future__ import print_function
-
-import torch
-import numpy as np
-from models.utils import Base
-from .priors import SplitPrior
-from .coupling import Coupling
-
-
-UNIT_TESTING = False
-
-
-def space_to_depth(x):
-    xs = x.size()
-    # Pick off every second element
-    x = x.view(xs[0], xs[1], xs[2] // 2, 2, xs[3] // 2, 2)
-    # Transpose picked elements next to channels.
-    x = x.permute((0, 1, 3, 5, 2, 4)).contiguous()
-    # Combine with channels.
-    x = x.view(xs[0], xs[1] * 4, xs[2] // 2, xs[3] // 2)
-    return x
-
-
-def depth_to_space(x):
-    xs = x.size()
-    # Pick off elements from channels
-    x = x.view(xs[0], xs[1] // 4, 2, 2, xs[2], xs[3])
-    # Transpose picked elements next to HW dimensions.
-    x = x.permute((0, 1, 4, 2, 5, 3)).contiguous()
-    # Combine with HW dimensions.
-    x = x.view(xs[0], xs[1] // 4, xs[2] * 2, xs[3] * 2)
-    return x
-
-
-def int_shape(x):
-    return list(map(int, x.size()))
-
-
-class Flatten(Base):
-    def forward(self, x):
-        return x.view(x.size(0), -1)
-
-
-class Reshape(Base):
-    def __init__(self, shape):
-        super().__init__()
-        self.shape = shape
-
-    def forward(self, x):
-        return x.view(x.size(0), *self.shape)
-
-
-class Reverse(Base):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, z, reverse=False):
-        flip_idx = torch.arange(z.size(1) - 1, -1, -1).long()
-        z = z[:, flip_idx, :, :]
-        return z
-
-
-class Permute(Base):
-    def __init__(self, n_channels):
-        super().__init__()
-
-        permutation = np.arange(n_channels, dtype='int')
-        np.random.shuffle(permutation)
-
-        permutation_inv = np.zeros(n_channels, dtype='int')
-        permutation_inv[permutation] = np.arange(n_channels, dtype='int')
-
-        self.permutation = torch.from_numpy(permutation)
-        self.permutation_inv = torch.from_numpy(permutation_inv)
-
-    def forward(self, z, ldj, reverse=False):
-        if not reverse:
-            z = z[:, self.permutation, :, :]
-        else:
-            z = z[:, self.permutation_inv, :, :]
-
-        return z, ldj
-
-    def InversePermute(self):
-        inv_permute = Permute(len(self.permutation))
-        inv_permute.permutation = self.permutation_inv
-        inv_permute.permutation_inv = self.permutation
-        return inv_permute
-
-
-class Squeeze(Base):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, z, ldj, reverse=False):
-        if not reverse:
-            z = space_to_depth(z)
-        else:
-            z = depth_to_space(z)
-        return z, ldj
-
-
-class GenerativeFlow(Base):
-    def __init__(self, n_channels, height, width, args):
-        super().__init__()
-        layers = []
-        layers.append(Squeeze())
-        n_channels *= 4
-        height //= 2
-        width //= 2
-
-        for level in range(args.n_levels):
-
-            for i in range(args.n_flows):
-                perm_layer = Permute(n_channels)
-                layers.append(perm_layer)
-
-                layers.append(
-                    Coupling(n_channels, height, width, args))
-
-            if level < args.n_levels - 1:
-                if args.splitprior_type != 'none':
-                    # Standard splitprior
-                    factor_out = n_channels // 2
-                    layers.append(SplitPrior(n_channels, factor_out, height, width, args))
-                    n_channels = n_channels - factor_out
-
-                layers.append(Squeeze())
-                n_channels *= 4
-                height //= 2
-                width //= 2
-
-        self.layers = torch.nn.ModuleList(layers)
-        self.z_size = (n_channels, height, width)
-
-    def forward(self, z, ldj, pys=(), ys=(), reverse=False):
-        if not reverse:
-            for l, layer in enumerate(self.layers):
-                if isinstance(layer, (SplitPrior)):
-                    py, y, z, ldj = layer(z, ldj)
-                    pys += (py,)
-                    ys += (y,)
-
-                else:
-                    z, ldj = layer(z, ldj)
-
-        else:
-            for l, layer in reversed(list(enumerate(self.layers))):
-                if isinstance(layer, (SplitPrior)):
-                    if len(ys) > 0:
-                        z, ldj = layer.inverse(z, ldj, y=ys[-1])
-                        # Pop last element
-                        ys = ys[:-1]
-                    else:
-                        z, ldj = layer.inverse(z, ldj, y=None)
-
-                else:
-                    z, ldj = layer(z, ldj, reverse=True)
-
-        return z, ldj, pys, ys
-
-    def decode(self, z, ldj, state, decode_fn):
-
-        for l, layer in reversed(list(enumerate(self.layers))):
-            if isinstance(layer, SplitPrior):
-                z, ldj, state = layer.decode(z, ldj, state, decode_fn)
-
-            else:
-                z, ldj = layer(z, ldj, reverse=True)
-
-        return z, ldj

integer_discrete_flows/models/networks.py

@@ -1,154 +0,0 @@
-"""
-Collection of flow strategies
-"""
-
-from __future__ import print_function
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from models.utils import Base
-
-
-UNIT_TESTING = False
-
-
-class Conv2dReLU(Base):
-    def __init__(
-            self, n_inputs, n_outputs, kernel_size=3, stride=1, padding=0,
-            bias=True):
-        super().__init__()
-
-        self.nn = nn.Conv2d(n_inputs, n_outputs, kernel_size, padding=padding)
-
-    def forward(self, x):
-        h = self.nn(x)
-
-        y = F.relu(h)
-
-        return y
-
-
-class ResidualBlock(Base):
-    def __init__(self, n_channels, kernel, Conv2dAct):
-        super().__init__()
-
-        self.nn = torch.nn.Sequential(
-            Conv2dAct(n_channels, n_channels, kernel, padding=1),
-            torch.nn.Conv2d(n_channels, n_channels, kernel, padding=1),
-            )
-
-    def forward(self, x):
-        h = self.nn(x)
-        h = F.relu(h + x)
-        return h
-
-
-class DenseLayer(Base):
-    def __init__(self, args, n_inputs, growth, Conv2dAct):
-        super().__init__()
-
-        conv1x1 = Conv2dAct(
-                n_inputs, n_inputs, kernel_size=1, stride=1,
-                padding=0, bias=True)
-
-        self.nn = torch.nn.Sequential(
-            conv1x1,
-            Conv2dAct(
-                n_inputs, growth, kernel_size=3, stride=1,
-                padding=1, bias=True),
-            )
-
-    def forward(self, x):
-        h = self.nn(x)
-
-        h = torch.cat([x, h], dim=1)
-        return h
-
-
-class DenseBlock(Base):
-    def __init__(
-            self, args, n_inputs, n_outputs, kernel, Conv2dAct):
-        super().__init__()
-        depth = args.densenet_depth
-
-        future_growth = n_outputs - n_inputs
-
-        layers = []
-
-        for d in range(depth):
-            growth = future_growth // (depth - d)
-
-            layers.append(DenseLayer(args, n_inputs, growth, Conv2dAct))
-            n_inputs += growth
-            future_growth -= growth
-
-        self.nn = torch.nn.Sequential(*layers)
-
-    def forward(self, x):
-        return self.nn(x)
-
-
-class Identity(Base):
-    def __init__(self):
-        super.__init__()
-
-    def forward(self, x):
-        return x
-
-
-class NN(Base):
-    def __init__(
-            self, args, c_in, c_out, height, width, nn_type, kernel=3):
-        super().__init__()
-
-        Conv2dAct = Conv2dReLU
-        n_channels = args.n_channels
-
-        if nn_type == 'shallow':
-
-            if args.network1x1 == 'standard':
-                conv1x1 = Conv2dAct(
-                    n_channels, n_channels, kernel_size=1, stride=1,
-                    padding=0, bias=False)
-
-            layers = [
-                Conv2dAct(c_in, n_channels, kernel, padding=1),
-                conv1x1]
-
-            layers += [torch.nn.Conv2d(n_channels, c_out, kernel, padding=1)]
-
-        elif nn_type == 'resnet':
-            layers = [
-                Conv2dAct(c_in, n_channels, kernel, padding=1),
-                ResidualBlock(n_channels, kernel, Conv2dAct),
-                ResidualBlock(n_channels, kernel, Conv2dAct)]
-
-            layers += [
-                torch.nn.Conv2d(n_channels, c_out, kernel, padding=1)
-                ]
-
-        elif nn_type == 'densenet':
-            layers = [
-                DenseBlock(
-                    args=args,
-                    n_inputs=c_in,
-                    n_outputs=n_channels + c_in,
-                    kernel=kernel,
-                    Conv2dAct=Conv2dAct)]
-
-            layers += [
-                torch.nn.Conv2d(n_channels + c_in, c_out, kernel, padding=1)
-                ]
-        else:
-            raise ValueError
-
-        self.nn = torch.nn.Sequential(*layers)
-
-        # Set parameters of last conv-layer to zero.
-        if not UNIT_TESTING:
-            self.nn[-1].weight.data.zero_()
-            self.nn[-1].bias.data.zero_()
-
-    def forward(self, x):
-        return self.nn(x)

integer_discrete_flows/models/priors.py

@@ -1,164 +0,0 @@
-"""
-Collection of flow strategies
-"""
-
-from __future__ import print_function
-
-import torch
-import torch.nn.functional as F
-from torch.nn import Parameter
-from utils.distributions import sample_discretized_logistic, \
-    sample_mixture_discretized_logistic, sample_normal, sample_logistic, \
-    sample_discretized_normal, sample_mixture_normal
-from models.utils import Base
-from .networks import NN
-
-
-def sample_prior(px, variable_type, distribution_type, inverse_bin_width):
-    if variable_type == 'discrete':
-        if distribution_type == 'logistic':
-            if len(px) == 2:
-                return sample_discretized_logistic(
-                    *px, inverse_bin_width=inverse_bin_width)
-            elif len(px) == 3:
-                return sample_mixture_discretized_logistic(
-                    *px, inverse_bin_width=inverse_bin_width)
-
-        elif distribution_type == 'normal':
-            return sample_discretized_normal(
-                *px, inverse_bin_width=inverse_bin_width)
-
-    elif variable_type == 'continuous':
-        if distribution_type == 'logistic':
-            return sample_logistic(*px)
-        elif distribution_type == 'normal':
-            if len(px) == 2:
-                return sample_normal(*px)
-            elif len(px) == 3:
-                return sample_mixture_normal(*px)
-        elif distribution_type == 'steplogistic':
-            return sample_logistic(*px)
-
-    raise ValueError
-
-
-class Prior(Base):
-    def __init__(self, size, args):
-        super().__init__()
-        c, h, w = size
-
-        self.inverse_bin_width = 2**args.n_bits
-        self.variable_type = args.variable_type
-        self.distribution_type = args.distribution_type
-        self.n_mixtures = args.n_mixtures
-
-        if self.n_mixtures == 1:
-            self.mu = Parameter(torch.Tensor(c, h, w))
-            self.logs = Parameter(torch.Tensor(c, h, w))
-        elif self.n_mixtures > 1:
-            self.mu = Parameter(torch.Tensor(c, h, w, self.n_mixtures))
-            self.logs = Parameter(torch.Tensor(c, h, w, self.n_mixtures))
-            self.pi_logit = Parameter(torch.Tensor(c, h, w, self.n_mixtures))
-
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        self.mu.data.zero_()
-
-        if self.n_mixtures > 1:
-            self.pi_logit.data.zero_()
-            for i in range(self.n_mixtures):
-                self.mu.data[..., i] += i - (self.n_mixtures - 1) / 2.
-
-        self.logs.data.zero_()
-
-    def get_pz(self, n):
-        if self.n_mixtures == 1:
-            mu = self.mu.repeat(n, 1, 1, 1)
-            logs = self.logs.repeat(n, 1, 1, 1)  # scaling scale
-            return mu, logs
-
-        elif self.n_mixtures > 1:
-            pi = F.softmax(self.pi_logit, dim=-1)
-            mu = self.mu.repeat(n, 1, 1, 1, 1)
-            logs = self.logs.repeat(n, 1, 1, 1, 1)
-            pi = pi.repeat(n, 1, 1, 1, 1)
-            return mu, logs, pi
-
-    def forward(self, z, ldj):
-        pz = self.get_pz(z.size(0))
-
-        return pz, z, ldj
-
-    def sample(self, n):
-        pz = self.get_pz(n)
-
-        z_sample = sample_prior(pz, self.variable_type, self.distribution_type, self.inverse_bin_width)
-
-        return z_sample
-
-    def decode(self, states, decode_fn):
-        pz = self.get_pz(n=len(states))
-
-        states, z = decode_fn(states, pz)
-        return states, z
-
-
-class SplitPrior(Base):
-    def __init__(self, c_in, factor_out, height, width, args):
-        super().__init__()
-
-        self.split_idx = c_in - factor_out
-        self.inverse_bin_width = 2**args.n_bits
-        self.variable_type = args.variable_type
-        self.distribution_type = args.distribution_type
-        self.input_channel = c_in
-
-        self.nn = NN(
-            args=args,
-            c_in=c_in - factor_out,
-            c_out=factor_out * 2,
-            height=height,
-            width=width,
-            nn_type=args.splitprior_type)
-
-    def get_py(self, z):
-        h = self.nn(z)
-        mu = h[:, ::2, :, :]
-        logs = h[:, 1::2, :, :]
-
-        py = [mu, logs]
-
-        return py
-
-    def split(self, z):
-        z1 = z[:, :self.split_idx, :, :]
-        y = z[:, self.split_idx:, :, :]
-        return z1, y
-
-    def combine(self, z, y):
-        result = torch.cat([z, y], dim=1)
-
-        return result
-
-    def forward(self, z, ldj):
-        z, y = self.split(z)
-
-        py = self.get_py(z)
-
-        return py, y, z, ldj
-
-    def inverse(self, z, ldj, y):
-        # Sample if y is not given.
-        if y is None:
-            py = self.get_py(z)
-            y = sample_prior(py, self.variable_type, self.distribution_type, self.inverse_bin_width)
-
-        z = self.combine(z, y)
-
-        return z, ldj
-
-    def decode(self, z, ldj, states, decode_fn):
-        py = self.get_py(z)
-        states, y = decode_fn(states, py)
-        return self.combine(z, y), ldj, states

integer_discrete_flows/models/utils.py

@@ -1,36 +0,0 @@
-import torch
-
-
-class Base(torch.nn.Module):
-    """
-    The base class for modules. That contains a disable round mode
-    """
-
-    def __init__(self):
-        super().__init__()
-
-    def _set_child_attribute(self, attr, value):
-        r"""Sets the module in rounding mode.
-
-        This has any effect only on certain modules if variable type is
-        discrete.
-
-        Returns:
-            Module: self
-        """
-        if hasattr(self, attr):
-            setattr(self, attr, value)
-
-        for module in self.modules():
-            if hasattr(module, attr):
-                setattr(module, attr, value)
-        return self
-
-    def set_temperature(self, value):
-        self._set_child_attribute("temperature", value)
-
-    def enable_hard_round(self, mode=True):
-        self._set_child_attribute("hard_round", mode)
-
-    def disable_hard_round(self, mode=True):
-        self.enable_hard_round(not mode)

integer_discrete_flows/optimization/loss.py

@@ -1,148 +0,0 @@
-from __future__ import print_function
-
-import numpy as np
-import torch
-from utils.distributions import log_discretized_logistic, \
-    log_mixture_discretized_logistic, log_normal, log_discretized_normal, \
-    log_logistic, log_mixture_normal
-from models.backround import _round_straightthrough
-
-
-def compute_log_ps(pxs, xs, args):
-    # Add likelihoods of intermediate representations.
-    inverse_bin_width = 2.**args.n_bits
-
-    log_pxs = []
-    for px, x in zip(pxs, xs):
-
-        if args.variable_type == 'discrete':
-            if args.distribution_type == 'logistic':
-                log_px = log_discretized_logistic(
-                    x, *px, inverse_bin_width=inverse_bin_width)
-            elif args.distribution_type == 'normal':
-                log_px = log_discretized_normal(
-                    x, *px, inverse_bin_width=inverse_bin_width)
-        elif args.variable_type == 'continuous':
-            if args.distribution_type == 'logistic':
-                log_px = log_logistic(x, *px)
-            elif args.distribution_type == 'normal':
-                log_px = log_normal(x, *px)
-            elif args.distribution_type == 'steplogistic':
-                x = _round_straightthrough(x * inverse_bin_width) / inverse_bin_width
-                log_px = log_discretized_logistic(
-                    x, *px, inverse_bin_width=inverse_bin_width)
-
-        log_pxs.append(
-            torch.sum(log_px, dim=[1, 2, 3]))
-
-    return log_pxs
-
-
-def compute_log_pz(pz, z, args):
-    inverse_bin_width = 2.**args.n_bits
-
-    if args.variable_type == 'discrete':
-        if args.distribution_type == 'logistic':
-            if args.n_mixtures == 1:
-                log_pz = log_discretized_logistic(
-                    z, pz[0], pz[1], inverse_bin_width=inverse_bin_width)
-            else:
-                log_pz = log_mixture_discretized_logistic(
-                    z, pz[0], pz[1], pz[2],
-                    inverse_bin_width=inverse_bin_width)
-        elif args.distribution_type == 'normal':
-            log_pz = log_discretized_normal(
-                z, *pz, inverse_bin_width=inverse_bin_width)
-
-    elif args.variable_type == 'continuous':
-        if args.distribution_type == 'logistic':
-            log_pz = log_logistic(z, *pz)
-        elif args.distribution_type == 'normal':
-            if args.n_mixtures == 1:
-                log_pz = log_normal(z, *pz)
-            else:
-                log_pz = log_mixture_normal(z, *pz)
-        elif args.distribution_type == 'steplogistic':
-            z = _round_straightthrough(z * 256.) / 256.
-            log_pz = log_discretized_logistic(z, *pz)
-
-    log_pz = torch.sum(
-        log_pz,
-        dim=[1, 2, 3])
-
-    return log_pz
-
-
-def compute_loss_function(pz, z, pys, ys, ldj, args):
-    """
-    Computes the cross entropy loss function while summing over batch dimension, not averaged!
-    :param x_logit: shape: (batch_size, num_classes * num_channels, pixel_width, pixel_height), real valued logits
-    :param x: shape (batchsize, num_channels, pixel_width, pixel_height), pixel values rescaled between [0, 1].
-    :param z_mu: mean of z_0
-    :param z_var: variance of z_0
-    :param z_0: first stochastic latent variable
-    :param z_k: last stochastic latent variable
-    :param ldj: log det jacobian
-    :param args: global parameter settings
-    :param beta: beta for kl loss
-    :return: loss, ce, kl
-    """
-    batch_size = z.size(0)
-
-    # Get array loss, sum over batch
-    loss_array, bpd_array, bpd_per_prior_array = \
-        compute_loss_array(pz, z, pys, ys, ldj, args)
-
-    loss = torch.mean(loss_array)
-    bpd = torch.mean(bpd_array).item()
-    bpd_per_prior = [torch.mean(x) for x in bpd_per_prior_array]
-
-    return loss, bpd, bpd_per_prior
-
-
-def convert_bpd(log_p, input_size):
-    return -log_p / (np.prod(input_size) * np.log(2.))
-
-
-def compute_loss_array(pz, z, pys, ys, ldj, args):
-    """
-    Computes the cross entropy loss function while summing over batch dimension, not averaged!
-    :param x_logit: shape: (batch_size, num_classes * num_channels, pixel_width, pixel_height), real valued logits
-    :param x: shape (batchsize, num_channels, pixel_width, pixel_height), pixel values rescaled between [0, 1].
-    :param z_mu: mean of z_0
-    :param z_var: variance of z_0
-    :param z_0: first stochastic latent variable
-    :param z_k: last stochastic latent variable
-    :param ldj: log det jacobian
-    :param args: global parameter settings
-    :param beta: beta for kl loss
-    :return: loss, ce, kl
-    """
-    bpd_per_prior = []
-
-    # Likelihood of final representation.
-    log_pz = compute_log_pz(pz, z, args)
-
-    bpd_per_prior.append(convert_bpd(log_pz.detach(), args.input_size))
-
-    log_p = log_pz
-
-    # Add likelihoods of intermediate representations.
-    if ys:
-        log_pys = compute_log_ps(pys, ys, args)
-
-        for log_py in log_pys:
-            log_p += log_py
-
-            bpd_per_prior.append(convert_bpd(log_py.detach(), args.input_size))
-
-    log_p += ldj
-
-    loss = -log_p
-    bpd = convert_bpd(log_p.detach(), args.input_size)
-
-    return loss, bpd, bpd_per_prior
-
-
-def calculate_loss(pz, z, pys, ys, ldj, loss_aux, args):
-    return compute_loss_function(pz, z, pys, ys, ldj, loss_aux, args)

integer_discrete_flows/optimization/training.py

@@ -1,174 +0,0 @@
-from __future__ import print_function
-import torch
-
-from optimization.loss import calculate_loss
-from utils.visual_evaluation import plot_reconstructions
-
-import numpy as np
-
-
-def train(epoch, train_loader, model, opt, args):
-    model.train()
-    train_loss = np.zeros(len(train_loader))
-    train_bpd = np.zeros(len(train_loader))
-
-    num_data = 0
-
-    for batch_idx, (data, _) in enumerate(train_loader):
-        data = data.view(-1, *args.input_size)
-
-        data = data.to(args.DEVICE)
-
-        opt.zero_grad()
-        loss, bpd, bpd_per_prior, pz, z, pys, py, ldj = model(data)
-
-        loss = torch.mean(loss)
-        bpd = torch.mean(bpd)
-        bpd_per_prior = [torch.mean(i) for i in bpd_per_prior]
-
-        loss.backward()
-        loss = loss.item()
-        train_loss[batch_idx] = loss
-        train_bpd[batch_idx] = bpd
-
-        ldj = torch.mean(ldj).item() / np.prod(args.input_size) / np.log(2)
-
-        opt.step()
-
-        num_data += len(data)
-
-        if batch_idx % args.log_interval == 0:
-            perc = 100. * batch_idx / len(train_loader)
-
-            tmp = 'Epoch: {:3d} [{:5d}/{:5d} ({:2.0f}%)] \tLoss: {:11.6f}\tbpd: {:8.6f}\tbits ldj: {:8.6f}'
-            print(tmp.format(epoch, num_data, len(train_loader.sampler), perc, loss, bpd, ldj))
-
-            print('z min: {:8.3f}, max: {:8.3f}'.format(torch.min(z).item() * 256, torch.max(z).item() * 256))
-
-            print('z  bpd: {:.3f}'.format(bpd_per_prior[0]))
-            for i in range(1, len(bpd_per_prior)):
-                print('y{} bpd: {:.3f}'.format(i-1, bpd_per_prior[i]))
-
-            print('pz mu', np.mean(pz[0].data.cpu().numpy(), axis=(0, 1, 2, 3)))
-            print('pz logs ', np.mean(pz[1].data.cpu().numpy(), axis=(0, 1, 2, 3)))
-            if len(pz) == 3:
-                print('pz pi   ', np.mean(pz[2].data.cpu().numpy(), axis=(0, 1, 2, 3)))
-
-            for i, py in enumerate(pys):
-                print('py{} mu   '.format(i), np.mean(py[0].data.cpu().numpy(), axis=(0, 1, 2, 3)))
-                print('py{} logs '.format(i), np.mean(py[1].data.cpu().numpy(), axis=(0, 1, 2, 3)))
-
-    from utils.visual_evaluation import plot_images
-    import os
-    if not os.path.exists(args.snap_dir + 'training/'):
-        os.makedirs(args.snap_dir + 'training/')
-
-    print('====> Epoch: {:3d} Average train loss: {:.4f}, average bpd: {:.4f}'.format(
-        epoch, train_loss.sum() / len(train_loader), train_bpd.sum() / len(train_loader)))
-
-    return train_loss, train_bpd
-
-
-def evaluate(train_loader, val_loader, model, model_sample, args, testing=False, file=None, epoch=0):
-    model.eval()
-
-    loss_type = 'bpd'
-
-    def analyse(data_loader, plot=False):
-        bpds = []
-        batch_idx = 0
-        with torch.no_grad():
-            for data, _ in data_loader:
-                batch_idx += 1
-
-                if args.cuda:
-                    data = data.cuda()
-
-                data = data.view(-1, *args.input_size)
-
-                loss, batch_bpd, bpd_per_prior, pz, z, pys, ys, ldj = \
-                    model(data)
-                loss = torch.mean(loss).item()
-                batch_bpd = torch.mean(batch_bpd).item()
-
-                bpds.append(batch_bpd)
-
-        bpd = np.mean(bpds)
-
-        with torch.no_grad():
-            if not testing and plot:
-                x_sample = model_sample.sample(n=100)
-
-                try:
-                    plot_reconstructions(
-                        x_sample, bpd, loss_type, epoch, args)
-                except:
-                    print('Not plotting')
-
-        return bpd
-
-    bpd_train = analyse(train_loader)
-    bpd_val = analyse(val_loader, plot=True)
-
-    with open(file, 'a') as ff:
-        msg = 'epoch {}\ttrain bpd {:.3f}\tval bpd {:.3f}\t'.format(
-                epoch,
-                bpd_train,
-                bpd_val)
-        print(msg, file=ff)
-
-    loss = bpd_val * np.prod(args.input_size) * np.log(2.)
-    bpd = bpd_val
-
-    file = None
-
-    # Compute log-likelihood
-    with torch.no_grad():
-        if testing:
-            test_data = val_loader.dataset.data_tensor
-
-            if args.cuda:
-                test_data = test_data.cuda()
-
-            print('Computing log-likelihood on test set')
-
-            model.eval()
-
-            log_likelihood = analyse(test_data)
-
-        else:
-            log_likelihood = None
-            nll_bpd = None
-
-        if file is None:
-            if testing:
-                print('====> Test set loss: {:.4f}'.format(loss))
-                print('====> Test set log-likelihood: {:.4f}'.format(log_likelihood))
-
-                print('====> Test set bpd (elbo): {:.4f}'.format(bpd))
-                print('====> Test set bpd (log-likelihood): {:.4f}'.format(log_likelihood/
-                                                                           (np.prod(args.input_size) * np.log(2.))))
-
-            else:
-                print('====> Validation set loss: {:.4f}'.format(loss))
-                print('====> Validation set bpd: {:.4f}'.format(bpd))
-        else:
-            with open(file, 'a') as ff:
-                if testing:
-                    print('====> Test set loss: {:.4f}'.format(loss), file=ff)
-                    print('====> Test set log-likelihood: {:.4f}'.format(log_likelihood), file=ff)
-
-                    print('====> Test set bpd: {:.4f}'.format(bpd), file=ff)
-                    print('====> Test set bpd (log-likelihood): {:.4f}'.format(log_likelihood /
-                                                                               (np.prod(args.input_size) * np.log(2.))),
-                          file=ff)
-
-                else:
-                    print('====> Validation set loss: {:.4f}'.format(loss), file=ff)
-                    print('====> Validation set bpd: {:.4f}'.format(loss / (np.prod(args.input_size) * np.log(2.))),
-                          file=ff)
-
-    if not testing:
-        return loss, bpd
-    else:
-        return log_likelihood, nll_bpd

integer_discrete_flows/utils/distributions.py

@@ -1,209 +0,0 @@
-from __future__ import print_function
-import torch
-import torch.utils.data
-import torch.nn.functional as F
-
-import numpy as np
-import math
-
-MIN_EPSILON = 1e-5
-MAX_EPSILON = 1.-1e-5
-
-
-PI = math.pi
-
-
-def log_min_exp(a, b, epsilon=1e-8):
-    """
-    Computes the log of exp(a) - exp(b) in a (more) numerically stable fashion.
-    Using:
-     log(exp(a) - exp(b))
-     c + log(exp(a-c) - exp(b-c))
-     a + log(1 - exp(b-a))
-    And note that we assume b < a always.
-    """
-    y = a + torch.log(1 - torch.exp(b - a) + epsilon)
-
-    return y
-
-
-def log_normal(x, mean, logvar):
-    logp = -0.5 * logvar
-    logp += -0.5 * np.log(2 * PI)
-    logp += -0.5 * (x - mean) * (x - mean) / torch.exp(logvar)
-    return logp
-
-
-def log_mixture_normal(x, mean, logvar, pi):
-    x = x.view(x.size(0), x.size(1), x.size(2), x.size(3), 1)
-
-    logp_mixtures = log_normal(x, mean, logvar)
-
-    logp = torch.log(torch.sum(pi * torch.exp(logp_mixtures), dim=-1) + 1e-8)
-
-    return logp
-
-
-def sample_normal(mean, logvar):
-    y = torch.randn_like(mean)
-
-    x = torch.exp(0.5 * logvar) * y + mean
-
-    return x
-
-
-def sample_mixture_normal(mean, logvar, pi):
-    b, c, h, w, n_mixtures = tuple(map(int, pi.size()))
-    pi = pi.view(b * c * h * w, n_mixtures)
-    sampled_pi = torch.multinomial(pi, num_samples=1).view(-1)
-
-    # Select mixture params
-    mean = mean.view(b * c * h * w, n_mixtures)
-    mean = mean[torch.arange(b*c*h*w), sampled_pi].view(b, c, h, w)
-    logvar = logvar.view(b * c * h * w, n_mixtures)
-    logvar = logvar[torch.arange(b*c*h*w), sampled_pi].view(b, c, h, w)
-
-    y = sample_normal(mean, logvar)
-
-    return y
-
-
-def log_logistic(x, mean, logscale):
-    """
-       pdf = sigma([x - mean] / scale) * [1 - sigma(...)] * 1/scale
-    """
-    scale = torch.exp(logscale)
-
-    u = (x - mean) / scale
-
-    logp = F.logsigmoid(u) + F.logsigmoid(-u) - logscale
-
-    return logp
-
-
-def sample_logistic(mean, logscale):
-    y = torch.rand_like(mean)
-
-    x = torch.exp(logscale) * torch.log(y / (1 - y)) + mean
-
-    return x
-
-
-def log_discretized_logistic(x, mean, logscale, inverse_bin_width):
-    scale = torch.exp(logscale)
-
-    logp = log_min_exp(
-        F.logsigmoid((x + 0.5 / inverse_bin_width - mean) / scale),
-        F.logsigmoid((x - 0.5 / inverse_bin_width - mean) / scale))
-
-    return logp
-
-
-def discretized_logistic_cdf(x, mean, logscale, inverse_bin_width):
-    scale = torch.exp(logscale)
-
-    cdf = torch.sigmoid((x + 0.5 / inverse_bin_width - mean) / scale)
-
-    return cdf
-
-
-def sample_discretized_logistic(mean, logscale, inverse_bin_width):
-    x = sample_logistic(mean, logscale)
-
-    x = torch.round(x * inverse_bin_width) / inverse_bin_width
-    return x
-
-
-def normal_cdf(value, loc, std):
-        return 0.5 * (1 + torch.erf((value - loc) * std.reciprocal() / math.sqrt(2)))
-
-
-def log_discretized_normal(x, mean, logvar, inverse_bin_width):
-    std = torch.exp(0.5 * logvar)
-    log_p = torch.log(normal_cdf(x + 0.5 / inverse_bin_width, mean, std) - normal_cdf(x - 0.5 / inverse_bin_width, mean, std) + 1e-7)
-
-    return log_p
-
-
-def log_mixture_discretized_normal(x, mean, logvar, pi, inverse_bin_width):
-    std = torch.exp(0.5 * logvar)
-
-    x = x.view(x.size(0), x.size(1), x.size(2), x.size(3), 1)
-
-    p = normal_cdf(x + 0.5 / inverse_bin_width, mean, std) - normal_cdf(x - 0.5 / inverse_bin_width, mean, std)
-
-    p = torch.sum(p * pi, dim=-1)
-
-    logp = torch.log(p + 1e-8)
-
-    return logp
-
-
-def sample_discretized_normal(mean, logvar, inverse_bin_width):
-    y = torch.randn_like(mean)
-
-    x = torch.exp(0.5 * logvar) * y + mean
-
-    x = torch.round(x * inverse_bin_width) / inverse_bin_width
-
-    return x
-
-
-def log_mixture_discretized_logistic(x, mean, logscale, pi, inverse_bin_width):
-    scale = torch.exp(logscale)
-
-    x = x.view(x.size(0), x.size(1), x.size(2), x.size(3), 1)
-
-    p = torch.sigmoid((x + 0.5 / inverse_bin_width - mean) / scale) \
-        - torch.sigmoid((x - 0.5 / inverse_bin_width - mean) / scale)
-
-    p = torch.sum(p * pi, dim=-1)
-
-    logp = torch.log(p + 1e-8)
-
-    return logp
-
-
-def mixture_discretized_logistic_cdf(x, mean, logscale, pi, inverse_bin_width):
-    scale = torch.exp(logscale)
-
-    x = x[..., None]
-
-    cdfs = torch.sigmoid((x + 0.5 / inverse_bin_width - mean) / scale)
-
-    cdf = torch.sum(cdfs * pi, dim=-1)
-
-    return cdf
-
-
-def sample_mixture_discretized_logistic(mean, logs, pi, inverse_bin_width):
-    # Sample mixtures
-    b, c, h, w, n_mixtures = tuple(map(int, pi.size()))
-    pi = pi.view(b * c * h * w, n_mixtures)
-    sampled_pi = torch.multinomial(pi, num_samples=1).view(-1)
-
-    # Select mixture params
-    mean = mean.view(b * c * h * w, n_mixtures)
-    mean = mean[torch.arange(b*c*h*w), sampled_pi].view(b, c, h, w)
-    logs = logs.view(b * c * h * w, n_mixtures)
-    logs = logs[torch.arange(b*c*h*w), sampled_pi].view(b, c, h, w)
-
-    y = torch.rand_like(mean)
-    x = torch.exp(logs) * torch.log(y / (1 - y)) + mean
-
-    x = torch.round(x * inverse_bin_width) / inverse_bin_width
-
-    return x
-
-
-def log_multinomial(logits, targets):
-    return -F.cross_entropy(logits, targets, reduction='none')
-
-
-def sample_multinomial(logits):
-    b, n_categories, c, h, w = logits.size()
-    logits = logits.permute(0, 2, 3, 4, 1)
-    p = F.softmax(logits, dim=-1)
-    p = p.view(b * c * h * w, n_categories)
-    x = torch.multinomial(p, num_samples=1).view(b, c, h, w)
-    return x

integer_discrete_flows/utils/load_data.py

@@ -1,264 +0,0 @@
-from __future__ import print_function
-
-import numbers
-
-import torch
-import torch.utils.data as data_utils
-import pickle
-from scipy.io import loadmat
-
-import numpy as np
-
-import os
-from torch.utils.data import Dataset
-import torchvision
-from torchvision import transforms
-from torchvision.transforms import functional as vf
-from torch.utils.data import ConcatDataset
-
-from PIL import Image
-
-import os
-import os.path
-from os.path import join
-import sys
-import tarfile
-
-
-class ToTensorNoNorm():
-    def __call__(self, X_i):
-        return torch.from_numpy(np.array(X_i, copy=False)).permute(2, 0, 1)
-
-
-class PadToMultiple(object):
-    def __init__(self, multiple, fill=0, padding_mode='constant'):
-        assert isinstance(multiple, numbers.Number)
-        assert isinstance(fill, (numbers.Number, str, tuple))
-        assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric']
-
-        self.multiple = multiple
-        self.fill = fill
-        self.padding_mode = padding_mode
-
-    def __call__(self, img):
-        """
-        Args:
-            img (PIL Image): Image to be padded.
-        Returns:
-            PIL Image: Padded image.
-        """
-        w, h = img.size
-        m = self.multiple
-        nw = (w // m + int((w % m) != 0)) * m
-        nh = (h // m + int((h % m) != 0)) * m
-        padw = nw - w
-        padh = nh - h
-
-        out = vf.pad(img, (0, 0, padw, padh), self.fill, self.padding_mode)
-        return out
-
-    def __repr__(self):
-        return self.__class__.__name__ + '(multiple={0}, fill={1}, padding_mode={2})'.\
-            format(self.mulitple, self.fill, self.padding_mode)
-
-
-class CustomTensorDataset(Dataset):
-    """Dataset wrapping tensors.
-
-    Each sample will be retrieved by indexing tensors along the first dimension.
-
-    Arguments:
-        *tensors (Tensor): tensors that have the same size of the first dimension.
-    """
-
-    def __init__(self, *tensors, transform=None):
-        assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors)
-        self.tensors = tensors
-        self.transform = transform
-
-    def __getitem__(self, index):
-        from PIL import Image
-
-        X, y = self.tensors
-        X_i, y_i, = X[index], y[index]
-
-        if self.transform:
-            X_i = self.transform(X_i)
-            X_i = torch.from_numpy(np.array(X_i, copy=False))
-            X_i = X_i.permute(2, 0, 1)
-
-        return X_i, y_i
-
-    def __len__(self):
-        return self.tensors[0].size(0)
-
-
-def load_cifar10(args, **kwargs):
-    # set args
-    args.input_size = [3, 32, 32]
-    args.input_type = 'continuous'
-    args.dynamic_binarization = False
-
-    from keras.datasets import cifar10
-    (x_train, y_train), (x_test, y_test) = cifar10.load_data()
-
-    x_train = x_train.transpose(0, 3, 1, 2)
-    x_test = x_test.transpose(0, 3, 1, 2)
-
-    import math
-
-    if args.data_augmentation_level == 2:
-        data_transform = transforms.Compose([
-                transforms.ToPILImage(),
-                transforms.RandomHorizontalFlip(),
-                transforms.Pad(int(math.ceil(32 * 0.05)), padding_mode='edge'),
-                transforms.RandomAffine(degrees=0, translate=(0.05, 0.05)),
-                transforms.CenterCrop(32)
-            ])
-    elif args.data_augmentation_level == 1:
-        data_transform = transforms.Compose([
-                transforms.ToPILImage(),
-                transforms.RandomHorizontalFlip(),
-            ])
-    else:
-        data_transform = transforms.Compose([
-                transforms.ToPILImage(),
-            ])
-
-    x_val = x_train[-10000:]
-    y_val = y_train[-10000:]
-
-    x_train = x_train[:-10000]
-    y_train = y_train[:-10000]
-
-    train = CustomTensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train), transform=data_transform)
-    train_loader = data_utils.DataLoader(train, batch_size=args.batch_size, shuffle=True, **kwargs)
-
-    validation = data_utils.TensorDataset(torch.from_numpy(x_val), torch.from_numpy(y_val))
-    val_loader = data_utils.DataLoader(validation, batch_size=args.batch_size, shuffle=False, **kwargs)
-
-    test = data_utils.TensorDataset(torch.from_numpy(x_test), torch.from_numpy(y_test))
-    test_loader = data_utils.DataLoader(test, batch_size=args.batch_size, shuffle=False, **kwargs)
-
-    return train_loader, val_loader, test_loader, args
-
-
-def extract_tar(tarpath):
-    assert tarpath.endswith('.tar')
-
-    startdir = tarpath[:-4] + '/'
-
-    if os.path.exists(startdir):
-        return startdir
-
-    print('Extracting', tarpath)
-
-    with tarfile.open(name=tarpath) as tar:
-        t = 0
-        done = False
-        while not done:
-            path = join(startdir, 'images{}'.format(t))
-            os.makedirs(path, exist_ok=True)
-
-            print(path)
-
-            for i in range(50000):
-                member = tar.next()
-
-                if member is None:
-                    done = True
-                    break
-
-                # Skip directories
-                while member.isdir():
-                    member = tar.next()
-                    if member is None:
-                        done = True
-                        break
-
-                member.name = member.name.split('/')[-1]
-
-                tar.extract(member, path=path)
-
-            t += 1
-
-    return startdir
-
-
-def load_imagenet(resolution, args, **kwargs):
-    assert resolution == 32 or resolution == 64
-
-    args.input_size = [3, resolution, resolution]
-
-    trainpath = '../imagenet{res}/train_{res}x{res}.tar'.format(res=resolution)
-    valpath = '../imagenet{res}/valid_{res}x{res}.tar'.format(res=resolution)
-
-    trainpath = extract_tar(trainpath)
-    valpath = extract_tar(valpath)
-
-    data_transform = transforms.Compose([
-        ToTensorNoNorm()
-    ])
-
-    print('Starting loading ImageNet')
-
-    imagenet_data = torchvision.datasets.ImageFolder(
-        trainpath,
-        transform=data_transform)
-
-    print('Number of data images', len(imagenet_data))
-
-    val_idcs = np.random.choice(len(imagenet_data), size=20000, replace=False)
-    train_idcs = np.setdiff1d(np.arange(len(imagenet_data)), val_idcs)
-
-    train_dataset = torch.utils.data.dataset.Subset(
-        imagenet_data, train_idcs)
-    val_dataset = torch.utils.data.dataset.Subset(
-        imagenet_data, val_idcs)
-
-    train_loader = torch.utils.data.DataLoader(
-        train_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        **kwargs)
-
-    val_loader = torch.utils.data.DataLoader(
-        val_dataset,
-        batch_size=args.batch_size,
-        shuffle=False,
-        **kwargs)
-
-    test_dataset = torchvision.datasets.ImageFolder(
-        valpath,
-        transform=data_transform)
-
-    print('Number of val images:', len(test_dataset))
-
-    test_loader = torch.utils.data.DataLoader(
-        test_dataset,
-        batch_size=args.batch_size,
-        shuffle=False,
-        **kwargs)
-
-    return train_loader, val_loader, test_loader, args
-
-
-def load_dataset(args, **kwargs):
-
-    if args.dataset == 'cifar10':
-        train_loader, val_loader, test_loader, args = load_cifar10(args, **kwargs)
-    elif args.dataset == 'imagenet32':
-        train_loader, val_loader, test_loader, args = load_imagenet(32, args, **kwargs)
-    elif args.dataset == 'imagenet64':
-        train_loader, val_loader, test_loader, args = load_imagenet(64, args, **kwargs)
-    else:
-        raise Exception('Wrong name of the dataset!')
-
-    return train_loader, val_loader, test_loader, args
-
-
-if __name__ == '__main__':
-    class Args():
-        def __init__(self):
-            self.batch_size = 128
-    train_loader, val_loader, test_loader, args = load_imagenet32(Args())

integer_discrete_flows/utils/log_likelihood.py

@@ -1,56 +0,0 @@
-from __future__ import print_function
-import numpy as np
-from scipy.misc import logsumexp
-from optimization.loss import calculate_loss_array
-
-
-def calculate_likelihood(X, model, args, S=5000, MB=500):
-
-    # set auxiliary variables for number of training and test sets
-    N_test = X.size(0)
-
-    X = X.view(-1, *args.input_size)
-
-    likelihood_test = []
-
-    if S <= MB:
-        R = 1
-    else:
-        R = S // MB
-        S = MB
-
-    for j in range(N_test):
-        if j % 100 == 0:
-            print('Progress: {:.2f}%'.format(j / (1. * N_test) * 100))
-
-        x_single = X[j].unsqueeze(0)
-
-        a = []
-        for r in range(0, R):
-            # Repeat it for all training points
-            x = x_single.expand(S, *x_single.size()[1:]).contiguous()
-
-            x_mean, z_mu, z_var, ldj, z0, zk = model(x)
-
-            a_tmp = calculate_loss_array(x_mean, x, z_mu, z_var, z0, zk, ldj, args)
-
-            a.append(-a_tmp.cpu().data.numpy())
-
-        # calculate max
-        a = np.asarray(a)
-        a = np.reshape(a, (a.shape[0] * a.shape[1], 1))
-        likelihood_x = logsumexp(a)
-        likelihood_test.append(likelihood_x - np.log(len(a)))
-
-    likelihood_test = np.array(likelihood_test)
-
-    nll = -np.mean(likelihood_test)
-
-    if args.input_type == 'multinomial':
-        bpd = nll/(np.prod(args.input_size) * np.log(2.))
-    elif args.input_type == 'binary':
-        bpd = 0.
-    else:
-        raise ValueError('invalid input type!')
-
-    return nll, bpd

integer_discrete_flows/utils/plotting.py

@@ -1,104 +0,0 @@
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-import matplotlib
-# noninteractive background
-matplotlib.use('Agg')
-import matplotlib.pyplot as plt
-
-
-def plot_training_curve(train_loss, validation_loss, fname='training_curve.pdf', labels=None):
-    """
-    Plots train_loss and validation loss as a function of optimization iteration
-    :param train_loss: np.array of train_loss (1D or 2D)
-    :param validation_loss: np.array of validation loss (1D or 2D)
-    :param fname: output file name
-    :param labels: if train_loss and validation loss are 2D, then labels indicate which variable is varied
-    accross training curves.
-    :return: None
-    """
-
-    plt.close()
-
-    matplotlib.rcParams.update({'font.size': 14})
-    matplotlib.rcParams['mathtext.fontset'] = 'stix'
-    matplotlib.rcParams['font.family'] = 'STIXGeneral'
-
-    if len(train_loss.shape) == 1:
-        # Single training curve
-        fig, ax = plt.subplots(nrows=1, ncols=1)
-        figsize = (6, 4)
-
-        if train_loss.shape[0] == validation_loss.shape[0]:
-            # validation score evaluated every iteration
-            x = np.arange(train_loss.shape[0])
-            ax.plot(x, train_loss, '-', lw=2., color='black', label='train')
-            ax.plot(x, validation_loss, '-', lw=2., color='blue', label='val')
-
-        elif train_loss.shape[0] % validation_loss.shape[0] == 0:
-            # validation score evaluated every epoch
-            x = np.arange(train_loss.shape[0])
-            ax.plot(x, train_loss, '-', lw=2., color='black', label='train')
-
-            x = np.arange(validation_loss.shape[0])
-            x = (x + 1) * train_loss.shape[0] / validation_loss.shape[0]
-            ax.plot(x, validation_loss, '-', lw=2., color='blue', label='val')
-        else:
-            raise ValueError('Length of train_loss and validation_loss must be equal or divisible')
-
-        miny = np.minimum(validation_loss.min(), train_loss.min()) - 20.
-        maxy = np.maximum(validation_loss.max(), train_loss.max()) + 30.
-        ax.set_ylim([miny, maxy])
-
-    elif len(train_loss.shape) == 2:
-        # Multiple training curves
-
-        cmap = plt.cm.brg
-
-        cNorm = matplotlib.colors.Normalize(vmin=0, vmax=train_loss.shape[0])
-        scalarMap = matplotlib.cm.ScalarMappable(norm=cNorm, cmap=cmap)
-
-        fig, ax = plt.subplots(nrows=1, ncols=1)
-        figsize = (6, 4)
-
-        if labels is None:
-            labels = ['%d' % i for i in range(train_loss.shape[0])]
-
-        if train_loss.shape[1] == validation_loss.shape[1]:
-            for i in range(train_loss.shape[0]):
-                color_val = scalarMap.to_rgba(i)
-
-                # validation score evaluated every iteration
-                x = np.arange(train_loss.shape[0])
-                ax.plot(x, train_loss[i], '-', lw=2., color=color_val, label=labels[i])
-                ax.plot(x, validation_loss[i], '--', lw=2., color=color_val)
-
-        elif train_loss.shape[1] % validation_loss.shape[1] == 0:
-            for i in range(train_loss.shape[0]):
-                color_val = scalarMap.to_rgba(i)
-
-                # validation score evaluated every epoch
-                x = np.arange(train_loss.shape[1])
-                ax.plot(x, train_loss[i], '-', lw=2., color=color_val, label=labels[i])
-
-                x = np.arange(validation_loss.shape[1])
-                x = (x+1) * train_loss.shape[1] / validation_loss.shape[1]
-                ax.plot(x, validation_loss[i], '-', lw=2., color=color_val)
-
-        miny = np.minimum(validation_loss.min(), train_loss.min()) - 20.
-        maxy = np.maximum(validation_loss.max(), train_loss.max()) + 30.
-        ax.set_ylim([miny, maxy])
-
-    else:
-        raise ValueError('train_loss and validation_loss must be 1D or 2D arrays')
-
-    ax.set_xlabel('iteration')
-    ax.set_ylabel('loss')
-    plt.title('Training and validation loss')
-
-    fig.set_size_inches(figsize)
-    fig.subplots_adjust(hspace=0.1)
-    plt.savefig(fname, bbox_inches='tight')
-
-    plt.close()

integer_discrete_flows/utils/visual_evaluation.py

@@ -1,37 +0,0 @@
-from __future__ import print_function
-import os
-import numpy as np
-import imageio
-
-
-def plot_reconstructions(recon_mean, loss, loss_type, epoch, args):
-    if epoch == 1:
-        if not os.path.exists(args.snap_dir + 'reconstruction/'):
-            os.makedirs(args.snap_dir + 'reconstruction/')
-    if loss_type == 'bpd':
-        fname = str(epoch) + '_bpd_%5.3f' % loss
-    elif loss_type == 'elbo':
-        fname = str(epoch) + '_elbo_%6.4f' % loss
-    plot_images(args, recon_mean.data.cpu().numpy()[:100], args.snap_dir + 'reconstruction/', fname)
-
-
-def plot_images(args, x_sample, dir, file_name, size_x=10, size_y=10):
-    batch, channels, height, width = x_sample.shape
-
-    print(x_sample.shape)
-
-    mosaic = np.zeros((height * size_y, width * size_x, channels))
-
-    for j in range(size_y):
-        for i in range(size_x):
-            idx = j * size_x + i
-
-            image = x_sample[idx]
-
-            mosaic[j*height:(j+1)*height, i*height:(i+1)*height] = \
-                image.transpose(1, 2, 0)
-
-    # Remove channel for BW images
-    mosaic = mosaic.squeeze()
-
-    imageio.imwrite(dir + file_name + '.png', mosaic)

main.py

@@ -89,7 +89,7 @@ def main():
 
     model = None
     if args.model_path is not None:
-        print("Loading the model...")
+        print("Loading the models...")
         model = torch.load(args.model_path)
 
     trainer: Trainer = OptunaTrainer(n_trials=3 if args.debug else None) if args.method == "optuna" else FullTrainer()

models/cnn/__init__.py

@@ -0,0 +1 @@
+from .cnn import CNNPredictor

trainers/FullTrainer.py

@@ -6,7 +6,7 @@ from torch.utils.data import DataLoader
 
 from .trainer import Trainer
 from .train import train
-from utils import print_losses
+from ..utils import print_losses
 
 class FullTrainer(Trainer):
     def execute(

trainers/OptunaTrainer.py

@@ -7,7 +7,7 @@ from torch import nn as nn
 from torch.utils.data import DataLoader
 
 from .trainer import Trainer
-from model.cnn import CNNPredictor
+from ..models.cnn import CNNPredictor
 from .train import train
 
 
@@ -59,4 +59,4 @@ class OptunaTrainer(Trainer):
         best_model = CNNPredictor(
             **best_params
         )
-        torch.save(best_model, "models/final_model.pt")
+        torch.save(best_model, f"saved_models/{model.__class__.__name__}.pt")