Add submitted code.

README.md

+# LNets
+Implementation and evaluation of Lipschitz neural networks (LNets). 
+
+# Installation
+* Create a new conda environment and activate it:
+    ```
+    conda create -n lnets python=3.5
+    conda activate lnets
+    ```
+    
+* Install PyTorch, following instructions in `https://pytorch.org`. 
+
+* Install torchnet by:
+    ```
+    pip install git+https://github.com/pytorch/tnt.git@master
+    ```
+
+* Navigate to the root of the project. Install the package, along with requirements:
+    ```
+    python setup.py install
+    ```
+    
+* Add project root to PYTHONPATH. One way to do this: 
+    ```
+    export PYTHONPATH="${PYTHONPATH}:`pwd`"
+    ``` 
+    
+**Note on PyTorch version**: All the experiments were performed using PyTorch version 0.4, although the code is expected
+to run using Pytorch 1.0. 
+
+# Models
+Code that implements the core ideas presented in the paper are shown below. 
+
+```
+lnets
+├── models
+│   └── acivations
+│       └── group_sort.py                   "GroupSort activation. " 
+│       └── maxout.py                       "MaxOut and MaxMin activations. "
+│   └── layers
+│       └── conv
+│           └── bjorck_conv2d.py            "Conv layer with Bjork-orthonormalized filters. "
+│           └── l_inf_projected_conv2d.py   "Conv layer with L-infinity projected filters. "
+│       └── dense
+│           └── bjorck_linear.py            "Dense layer with Bjorck-orthonormalized weights. "
+│           └── l_inf_projected.py          "Dense layer with l-infinity projected weights. "
+│           └── parseval_l2_linear.py       "Dense layer with Parseval regularization. "
+│           └── spectral_normal.py          "Dense layer with spectral normalization. "
+│   └── regularization
+│       └── __init__.py
+│       └── spec_jac.py                     "Penalizes the jacobian norm. Description in Appendix of paper. "
+│   └── utils
+│       └── __init__.py
+│       └── conversion.py                   "Converts a Bjorck layer to a regular one for fast test time inference. "
+│   └── __init__.py                         "Specification of models for a variety of tasks. "
+```
+
+## Configuring Experiments
+We strived to put as many variables as we could in a single configuration (json) file for each experiment. 
+Sample configuration files exist under:
+* `lnets/tasks/adversarial/configs`: for adversarial robustness experiments.
+* `lnets/tasks/classification/configs`: for classification experimnts.
+* `lnets/tasks/dualnets/configs`: for Wasserstein distance estimation experiments. 
+* `lnets/tasks/gan/configs`: for training GANs. 
+
+We now describe the key moving parts in these configs and how to change them.  
+
+### Model Configuration
+`model.name`: (string) Chooses the overall architecture and the task/training objective. `lnets/models/__init__.py` contains 
+ the commonly used model names. Two examples are:
+* "dual_fc": Train a fully connected model, under the dual Wasserstein objective. 
+* "classify_fc": Train a fully connected classifier. 
+
+`model.activation`: (string) Activation used throughout the network. One of "maxmin", "group_sort", "maxout", "relu", "tahn", 
+"sigmoid" or "identity" (i.e. no activation). 
+
+`model.linear.type`: (string) Chooses which linear layer type is going to be used. If the model is fully connected, the available
+options are:
+* "standard": The usual linear layer. 
+* "bjorck": Bjorck orthonormalized - all singular values equal to one. 
+* "l_inf_projected": Weight matrices are projected to the L-infinity ball. 
+* "spectral_normal": Use spectral normalization - largest singular value set to 1. 
+* "parseval_l2": Parseval regularized linear transformation.
+
+If the architecture is fully convolutional, the available options are: 
+* "standard_conv2d": the standard convolutional layer,
+* "bjorck_conv2d": Convolutional layers in which the filters are Bjorck orthonormalized,
+* "l_inf_projected_conv2d": Convolutional layers in which the filters are projected to the L-infinity ball. 
+
+`model.layers`: (list) Contains how many neurons (or convolutional filters) there should be in each layer. 
+
+`model.groupings`: (list) This field is used for activations that perform operations on groups of neurons. Used for GroupSort, 
+MaxMin and MaxOut. Is a list specifying the grouping sizes for each layer. For example, setting to \[2, 3\] means the 
+activation should act on groups of 2 and 3 in the first and second layers, respectively.
+
+`l_constant`: (integer) Scales the output of each layer by a certain amount such that the network output is scaled by 
+l_constant. Used to build K-Lipschitz networks out of 1-Lipschitz building blocks. 
+
+`per_update_proj` and `per_epoch_proj`: Some algorithms (such as Parseval networks) involve projecting the weights of 
+networks to a certain manifold after each training update. These fields let the user flexibly choose how often and with
+which projection algorithm the weights should be projected. The supported projection algorithms are:
+* "l_2": project to L2 ball. 
+* "l_inf_projected": project to the L-infinity ball.
+ 
+ By default, after-update and after-epoch updates are set to false. 
+ 
+### Running on GPU
+If a GPU is available, we strongly encourage the users to turn on GPU training by turning on the related json field in
+the experiment configs. In all experiments, set  `"cuda": true` (except for the GAN experiments, for which set
+`"gpu_mode": true`)
+ turn on the "cuda" field in the configurations. This speeds up
+training models significantly - especially with Bjorck layers. 
+ 
+### Other Configurations
+**Configuring optimizer**: Adam, standard SGD, nesterov momentum and AggMo are supported. Since most of the fields 
+in the optimizer configurations are self-explanatory, we leave it for the user to make use of the existing optimizer 
+configurations pushed in this repo. 
+
+**Miscellaneous**: Other fields control other aspects of training, such as IO settings, enabling cuda, logging, 
+visualizing results etc. 
+
+Other task specific configs will be described below under their corresponding titles. 
+
+# Tasks
+Four tasks are explored: Wasserstein Distance estimation, adversarial robustness, GAN training and classification. 
+
+### Wasserstein Distance Estimation
+ **Configuring Distributions**: The `distrib1` and `distrib2` fields are intended to be used to configure the probability
+distributions that will be used in the Wasserstein Distance estimation tasks. Currently, configs for 
+`multi_spherical _shell` (a distribution consisting of multiple spherical shells living in high dimensions) and 
+`gan_sampler` (samples from the empirical and generator distribution of a GAN) exist. 
+ 
+ 
+#### Quantifying Expressivity using Synthetic Distributions
+By using synthetic distributions whose Wasserstein distance and its accompanying dual surface we can analytically
+compute, we can quantify how expressive a Lipschitz architecture is. The closer the architecture can approximate the 
+correct Wasserstein distance, the more expressive it is. 
+
+* Approximating Absolute Value
+
+```
+python ./lnets/tasks/dualnets/mains/train_dual.py ./lnets/tasks/dualnets/configs/absolute_value_experiment.json
+```
+
+* Approximating Three Cones
+
+```
+python ./lnets/tasks/dualnets/mains/train_dual.py ./lnets/tasks/dualnets/configs/three_cones_experiment.json
+```
+
+* Approximating High Dimensional Cones
+
+```
+python ./lnets/tasks/dualnets/mains/train_dual.py ./lnets/tasks/dualnets/configs/high_dimensional_cone_experiment.json
+```
+
+#### Wasserstein Distance between GAN Generator and Empirical Distributions
+First, we need to train a GAN so that we can use its generator network for the Wasserstein Distance estimation 
+task. 
+
+* GAN training 
+
+(defaults to training WGAN on MNIST)
+
+```
+python ./lnets/tasks/gan/mains/train_gan.py ./lnets/tasks/gan/configs/train_GAN.json
+```
+
+The GAN type and the training set (along with other training hyperparameters) can be changed:
+
+`gan_type`: One of "WGAN", "WGAN_GP" or "LWGAN" (where the discriminator consists of this paper's contributions - 
+more on this later)
+
+`dataset`: One of "mnist", "fashion-mnist", "cifar10", "svhn", "stl10" "lsun-bed"
+
+* Estimating Wasserstein Distance
+
+```
+python ./lnets/tasks/dualnets/mains/train_dual.py ./lnets/tasks/dualnets/configs/estimate_wde_gan.json           
+```
+
+In order to sample from the GAN trained in the above step, we need to modify the config used for wasserstein distance 
+estimation. 
+
+`distrib1.gan_config_json_path`: Path to the gan training config used in the first step. 
+
+One can then modify the model to see which Lipschitz architectures obtain a tighter lower bound on the Wasserstein
+distance between the generator and empirical data distribution. 
+
+### Training LWGAN (Lipschitz WGANs)
+We can use the same WGAN training methodology, but build a discriminator network comprised of our methods (i.e. Bjorck 
+orthonormalized linear transformations and GroupSort activations)
+
+#### Training LWGAN: 
+
+```
+python ./lnets/tasks/gan/mains/train_gan.py ./lnets/tasks/gan/configs/train_LWGAN.json
+```
+
+The the training set (along with other training hyperparameters) can be changed:
+
+`dataset`: One of "mnist", "fashion-mnist", "cifar10", "svhn", "stl10" "lsun-bed"
+
+### Classification
+
+#### Classification on Standard Datasets
+
+* Training a standard, fully connected classifier. 
+
+```
+python ./lnets/tasks/classification/mains/train_classifier.py ./lnets/tasks/classification/configs/standard/fc_classification.json
+```
+
+* Training Bjorck Lipschitz classifier
+
+```
+python ./lnets/tasks/classification/mains/train_classifier.py ./lnets/tasks/classification/configs/standard/fc_classification_bjorck.json -o model.linear.bjorck_iter=3
+```
+
+Note that we use few bjorck iterations for this training script. Lipschitz-ness will not be strictly enforced so
+we do additional finetuning afterwards.
+
+* Orthonormal finetuning
+
+```
+python ./lnets/tasks/classification/mains/ortho_finetune.py --model.exp_path=<trained_bjorck_model_path.pt>
+```
+
+#### Classification with Small Data
+* Generating data indices: Generate which samples in the dataset will be used for training. 
+
+```
+python ./lnets/tasks/classification/mains/generate_data_indices.py --data.name mnist --data.root "data/small_mnist" --data.class_count 10 --per_class_count 100 --val_size 5000
+```
+
+* Training small data classifier.
+
+```
+python ./lnets/tasks/classification/mains/train_classifier.py ./lnets/tasks/classification/configs/small_mnist/lenet_bjorck.json
+
+```
+
+### Adversarial Robustness
+
+For the robustness experiments we trained both the Bjorck orthonormal networks and the L-infinity max-margin networks.
+
+* Training  L-Inf Lipschitz margin network
+
+```
+python ./lnets/tasks/classification/mains/train_classifier.py ./lnets/tasks/classification/configs/standard/fc_classification_l_inf_margin.json
+```
+
+* Evaluating robustness of trained classifier
+
+```
+python ./lnets/tasks/adversarial/mains/manual_eval_adv_robustness.py --model.exp_path="root/of/above/experiment/results" --output_root="outs/adv_robustness/mnist_l_inf_margin"
+```
+
+## Code References
+
+* ResNet Implementation: Largely based on github/kuangliu - https://github.com/kuangliu/pytorch-cifar/blob/master/models/resnet.py
+* GAN training pipeline: Based on and refactored from https://github.com/znxlwm/pytorch-generative-model-collections

lnets/data/data_transforms.py

+import torchvision.transforms as transforms
+
+
+def get_data_transforms(config):
+    # train_transform = None
+    test_transform = None
+
+    if config.data.transform.type == 'cifar':
+        train_transform, test_transform = get_cifar_transform(config)
+    elif config.data.transform.type == 'imagenet':
+        train_transform, test_transform = get_imagenet_transform(config)
+    else:
+        train_transform = transforms.ToTensor()
+
+    # Make sure to turn the input images into PyTorch tensors.
+    if test_transform is None:
+        test_transform = transforms.ToTensor()
+
+    return train_transform, test_transform
+
+
+def get_cifar_transform(config):
+    normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],
+                                     std=[0.2023, 0.1994, 0.2010])
+    train_transform = transforms.Compose([
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        normalize
+    ])
+    test_transform = transforms.Compose([
+        transforms.ToTensor(),
+        normalize
+    ])
+
+    return train_transform, test_transform
+
+
+def get_imagenet_transform(config):
+    normalize = transforms.Normalize(mean=config.data.transform.norm_mean,
+                                     std=config.data.transform.norm_std)
+    train_transform = transforms.Compose([
+        transforms.RandomResizedCrop(224),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        normalize,
+    ])
+    test_transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        normalize,
+    ])
+
+    return train_transform, test_transform

lnets/data/generate_data_indices.py

+from lnets.data.load_data import get_datasets
+from lnets.data.utils import save_indices
+
+import argparse
+import os
+from munch import Munch
+
+
+def main(opt):
+    opt.data.transform = Munch(type='none')
+
+    indices_path = os.path.join(opt.data.root, opt.data.name)
+
+    train_data, _, _ = get_datasets(opt)
+
+    save_indices(train_data, indices_path, opt.per_class_count, opt.data.class_count, opt.val_size)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Generate data indices. ')
+    parser.add_argument('--data.name', type=str, metavar='MODELPATH',
+                        help="location of pretrained model weights to evaluate")
+    parser.add_argument('--data.root', type=str, help='output directory to which results should be saved')
+    parser.add_argument('--data.class_count', type=int, help='total number of classes in dataset')
+    parser.add_argument('--per_class_count', type=int, help="How many training data points per class")
+    parser.add_argument('--val_size', type=int, help="Total number of validation points")
+    args = vars(parser.parse_args())
+
+    opt = {}
+    for k, v in args.items():
+        cur = opt
+        tokens = k.split('.')
+        for token in tokens[:-1]:
+            if token not in cur:
+                cur[token] = {}
+            cur = cur[token]
+        cur[tokens[-1]] = v
+    main(Munch.fromDict(opt))

lnets/data/load_data.py

+import os
+from torch.utils.data import Subset, DataLoader
+import torchvision.datasets as datasets
+
+from lnets.data.data_transforms import get_data_transforms
+from lnets.data.utils import load_indices
+
+
+def get_datasets(config):
+    data_name = config['data']['name'].lower()
+    path = os.path.join(config['data']['root'], data_name)
+
+    train_transform, test_transform = get_data_transforms(config)
+
+    train_data_args = dict(download=True, transform=train_transform)
+    val_data_args = dict(download=True, transform=test_transform)
+    test_data_args = dict(train=False, download=True, transform=test_transform)
+
+    if data_name == 'mnist':
+        train_data = datasets.MNIST(path, **train_data_args)
+        val_data = datasets.MNIST(path, **val_data_args)
+        test_data = datasets.MNIST(path, **test_data_args)
+    elif data_name == 'cifar10':
+        train_data = datasets.CIFAR10(path, **train_data_args)
+        val_data = datasets.CIFAR10(path, **val_data_args)
+        test_data = datasets.CIFAR10(path, **test_data_args)
+    elif data_name == 'cifar100':
+        train_data = datasets.CIFAR100(path, **train_data_args)
+        val_data = datasets.CIFAR100(path, **val_data_args)
+        test_data = datasets.CIFAR100(path, **test_data_args)
+    elif data_name == 'fashion-mnist':
+        train_data = datasets.FashionMNIST(path, **train_data_args)
+        val_data = datasets.FashionMNIST(path, **val_data_args)
+        test_data = datasets.FashionMNIST(path, **test_data_args)
+    elif data_name == 'imagenet-torchvision':
+        train_data = datasets.ImageFolder(os.path.join(path, 'train'), transform=train_transform)
+        val_data = datasets.ImageFolder(os.path.join(path, 'valid'), transform=test_transform)
+        # Currently not loaded.
+        test_data = None
+    else:
+        raise NotImplementedError('Data name %s not supported' % data_name)
+
+    return train_data, val_data, test_data
+
+
+def build_loaders(config, train_data, val_data, test_data):
+    data_name = config['data']['name'].lower()
+    batch_size = config['optim']['batch_size']
+    num_workers = config['data']['num_workers']
+
+    if config['data']['indices_path'] is not None:
+        train_indices, val_indices = load_indices(config['data']['indices_path'], config['data']['per_class_count'])
+        train_data = Subset(train_data, train_indices)
+        val_data = Subset(val_data, val_indices)
+    elif data_name != 'imagenet-torchvision':
+        # Manually readjust train/val size for memory saving.
+        data_size = len(train_data)
+        train_size = int(data_size * config['data']['train_size'])
+
+        train_data.train_data = train_data.train_data[:train_size]
+        train_data.train_labels = train_data.train_labels[:train_size]
+
+        if config['data']['train_size'] != 1:
+            val_data.train_data = val_data.train_data[train_size:]
+            val_data.train_labels = val_data.train_labels[train_size:]
+        else:
+            val_data = None
+
+    loaders = {
+        'train': DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=num_workers),
+        'validation': DataLoader(val_data, batch_size=batch_size, num_workers=num_workers),
+        'test': DataLoader(test_data, batch_size=batch_size, num_workers=num_workers)
+    }
+
+    return loaders
+
+
+def load_data(config):
+    train_data, val_data, test_data = get_datasets(config)
+    return build_loaders(config, train_data, val_data, test_data)

lnets/data/small_data.py

+import numpy as np
+
+
+def get_small_data_indices(dataset, total_per_class, class_count, val_size):
+    total_points = len(dataset)
+    if total_per_class * class_count + val_size > total_points:
+        raise Exception('More data points requested than is in data')
+    random_indices = np.random.permutation(total_points)
+
+    small_data_indices = {}
+    val_indices = []
+    for c in range(class_count):
+        small_data_indices[c] = []
+
+    for idx in random_indices:
+        _, y = dataset[idx]
+        y = int(y.item())
+        if len(small_data_indices[y]) < total_per_class:
+            small_data_indices[y].append(idx)
+        elif len(val_indices) < val_size:
+            val_indices.append(idx)
+        if all([len(small_data_indices[c]) == total_per_class for c in range(class_count)]):
+            if len(val_indices) == val_size:
+                break
+    if not all([len(small_data_indices[c]) == total_per_class for c in range(class_count)]):
+        raise Warning('Uneven class counts in small data indices')
+    return np.array([small_data_indices[c] for c in
+                     range(class_count)]).astype(np.int32).flatten(), np.array(val_indices).astype(np.int32)

lnets/data/utils.py

+from lnets.data.small_data import get_small_data_indices
+
+import numpy as np
+import os
+
+
+def save_indices(dataset, indices_path, per_class_count, total_class_count, val_size):
+    train_indices, val_indices = get_small_data_indices(dataset, per_class_count, total_class_count, val_size)
+    np.savetxt(os.path.join(indices_path, "train_indices_{}.txt".format(per_class_count)), train_indices)
+    np.savetxt(os.path.join(indices_path, "val_indices_{}.txt".format(per_class_count)), val_indices)
+
+
+def load_indices(path, per_class_count):
+    train_indices = os.path.join(path, "train_indices_{}.txt".format(per_class_count))
+    val_indices = os.path.join(path, "val_indices_{}.txt".format(per_class_count))
+    return np.loadtxt(train_indices, dtype=np.int32), np.loadtxt(val_indices, dtype=np.int32)

lnets/models/__init__.py

+from lnets.models.architectures import *
+from lnets.models.model_types import *
+from lnets.models.layers import *
+
+MODEL_REGISTRY = {}
+
+
+def register_model(model_name):
+    def decorator(f):
+        MODEL_REGISTRY[model_name] = f
+        return f
+
+    return decorator
+
+
+def get_model(config):
+    model_name = config['model']['name']
+    if model_name in MODEL_REGISTRY:
+        return MODEL_REGISTRY[model_name](config)
+    else:
+        raise ValueError("Unknown model {:s}".format(model_name))
+
+
+# Wasserstein Distance Estimation.
+@register_model('dual_fc')
+def load_fc_dual(config):
+    model = FCNet(config.model.layers, config.distrib1.dim, config.model.linear.type, config.model.activation,
+                  bias=config.model.linear.bias, config=config)
+    return DualOptimModel(model)
+
+
+@register_model("dual_fully_conv")
+def load_conv_dual(config):
+    model = FullyConv2D(config.distrib1.dim, config.model.channels, config.model.kernels, config.model.strides,
+                        linear_type=config.model.linear.type, activation=config.model.activation, config=config)
+    return DualOptimModel(model)
+
+
+# Classification.
+@register_model('classify_fc')
+def load_classify_fc(config):
+    model = FCNet(config.model.layers, config.data.input_dim, config.model.linear.type, config.model.activation,
+                  bias=config.model.linear.bias, config=config)
+    return ClassificationModel(model)
+
+
+@register_model('classify_fc_dropout')
+def load_classify_fc_dropout(config):
+    model = FCNet(config.model.layers, config.data.input_dim, config.model.linear.type, config.model.activation,
+                  bias=config.model.linear.bias, config=config, dropout=True)
+    return ClassificationModel(model)
+
+
+@register_model('classify_fc_spec_jac')
+def load_classify_fc_spec_jac(config):
+    model = FCNet(config.model.layers, config.data.input_dim, config.model.linear.type, config.model.activation,
+                  bias=config.model.linear.bias, config=config)
+    return JacSpecClassificationModel(model, config['model']['sn_reg'], config['cuda'])
+
+
+@register_model('classify_fc_margin')
+def load_classify_fc_margin(config):
+    model = FCNet(config.model.layers, config.data.input_dim, config.model.linear.type, config.model.activation,
+                  bias=config.model.linear.bias, config=config)
+    return MarginClassificationModel(model, config)
+
+
+@register_model('classify_fc_hinge')
+def load_classify_fc_hinge(config):
+    model = FCNet(config.model.layers, config.data.input_dim, config.model.linear.type, config.model.activation,
+                  bias=config.model.linear.bias, config=config)
+    return HingeLossClassificationModel(model, config)
+
+
+@register_model("lenet_classify")
+def load_lenet_classify(config):
+    model = LeNet(config.data.in_channels, config.model.output_dim, config.model.linear.type, config.model.activation, config.model.dropout_on,
+                  config=config)
+    return ClassificationModel(model)
+
+
+@register_model('resnet32')
+def CifarResNet32(config):
+    block_config = {
+        "num_blocks": [5, 5, 5],
+        "num_channels": [16, 32, 64],
+        "width": 1,
+        "pool_size": 8
+    }
+    return ClassificationModel(ResNet(BasicBlock, block_config, config['data']['class_count']))
+
+
+@register_model('wide-resnet32')
+def CifarWideResNet32(config):
+    block_config = {
+        "num_blocks": [5, 5, 5],
+        "num_channels": [16, 32, 64],
+        "width": 10,
+        "pool_size": 8
+    }
+    return ClassificationModel(ResNet(BasicBlock, block_config, config['data']['class_count']))

lnets/models/activations/__init__.py

+from lnets.models.activations.base_activation import Activation
+from lnets.models.activations.maxout import Maxout, MaxMin
+from lnets.models.activations.identity import Identity
+from lnets.models.activations.group_sort import GroupSort

lnets/models/activations/base_activation.py

+import torch.nn as nn
+
+
+class Activation(nn.Module):
+    def __init__(self):
+        super(Activation, self).__init__()
+
+    def forward(self, x):
+        raise NotImplementedError

lnets/models/activations/group_sort.py

+import numpy as np
+import torch.nn as nn
+
+
+class GroupSort(nn.Module):
+
+    def __init__(self, num_units, axis=-1):
+        super(GroupSort, self).__init__()
+        self.num_units = num_units
+        self.axis = axis