Repository for reproducing `Model-Based Robust Deep Learning`

Related tags

Deep Learningmbrdl
Overview

Model-Based Robust Deep Learning (MBRDL)

In this repository, we include the code necessary for reproducing the code used in Model-Based Robust Deep Learning. In particular, we include the code necessary for both training models of natural variation as well as the code needed to train classifiers using these learned models. A brief summary of the functionality provided in this repo is provided below in the table of contents. If you find this repository useful in your research, please consider citing:

@article{robey2020model,
  title={Model-Based Robust Deep Learning},
  author={Robey, Alexander and Hassani, Hamed and Pappas, George J},
  journal={arXiv preprint arXiv:2005.10247},
  year={2020}
}

Table of contents

First, we given instructions for how to setup the appropriate environment for this repository.

Next, we give details about how to train classifiers using the MBRDL paradigm. ur implementation is based on the Lambda Labs implementation of the ImageNet training repository.

Following this, we give a high-level overview of the structure of this repository, as well as an overview of the scripts that are included for training models of natural variation and classifiers that are robust to natural variation.

We also provide code that can be used to train models of natural variation using the MUNIT framework. The code that we use to train these models is largely based on the original implementation of MUNIT.

In addition to providing functionality to train new models of natural variation, we also provide a library of pre-trained models of natural variation in a Google Drive folder.

Finally, we identify common usage issues and provide solutions.

Setup instructions

After cloning this repository, the first step is to setup a virtual environment.

python3 -m venv mbrdl
source mbrdl/bin/activate
pip3 install -r requirements.txt
pip3 install https://download.pytorch.org/whl/cu100/torch-1.1.0-cp36-cp36m-linux_x86_64.whl

We also need to install NVIDIA's half-precision training tool apex. The setup instructions for apex are here.

An overview of the structure of this repository

In this repository, we provide four scripts that can be used to train models of natural variation as well as classifiers that are robust against natural variation; the scripts train_munit.sh and train_munit_imagenet.sh can be used to train models of natural variation with the MUNIT framework, and the scripts train_basic.sh and train_imagenet.sh can be used to train classifiers using the MBRDL paradigm.

All of the components to train models of natural variation and classifiers is included in the core/ directory. In particular, core/ is organized in the following way:

mbrdl/core/
├── classifiers
├── data
├── models
├── training
├── train_munit.py
├── train.py
└── utils

The classifiers/ directory contains all code necessary to initialize different classifier architectures. The data/ directory contains dataloaders for MUNIT and for training classifiers for SVHN, CURE-TSR, GTSRB, and ImageNet/ImageNet-c. The models/ directory has code that can be used to load and train models of natural variation using the MUNIT framework. The training/ directory contains the training algorithms (including implementations of MAT, MRT, and MDA) and schedulers that are used to train classifiers. The utils/ directory has a variety of utilities that are used throughout the scripts that train classifiers in the MBRDL paradigm.

The two main training python files are train_munit.py and train.py. As the name suggests, train_munit.py can be used to train models of natural variation using MUNIT. train.py is a python file that trains classifiers in a distributed fashion.

Training classifiers in the MBRDL paradigm

We also include the code needed to train classifiers that are robust to natural variation. Our implementation is based on the Lambda Labs implementation of the ImageNet training repository.

To train a classifier on MNIST, SVHN, GTSRB, or CURE-TSR, you can run the following shell script:

chmod +x train_basic.sh
./train_basic.sh

Editing the train_basic.sh script will allow you to control the dataset, optimization parameters, model of natural variation, classifier architecture, and other hyperparameters. To train on ImageNet/ImageNet-c, we provide a separate script, which can be run in the following way:

chmod +x train_imagenet_script.sh
./train_imagenet.sh

In what follows, we describe different settings for thes files.

Dataset selection

By editing flags at the beginning of the file, you can change the dataset and the source of natural variation that are used for training/testing the classifier. For example, to run with MNIST, you can set:

export DATASET='mnist'
export TRAIN_DIR=./datasets/MNIST/mnist
export SOURCE='background-color'

To run with SVHN, you can set:

export DATASET='svhn'
export TRAIN_DIR=./datasets/svhn
export SOURCE='brightness'

The choices for SOURCE for SVHN are choices are 'brightness', 'contrast', and 'contrast+brightness'. The same choices are also available for GTSRB:

export DATASET='gtsrb'
export TRIAIN_DIR=./datasets/gtsrb/GTSRB
export SOURCE='brightness'

For CURE-TSR, you can select any of the sources of natural variation listed in the original repository for CURE-TSR (e.g. snow, rain, haze, decolorization, etc.). For example, you can set

export DATASET='cure-tsr'
export TRAIN_DIR=./datasets/cure_tsr/raw_data
export SOURCE='snow'

To train with ImageNet, you need to set the TRAIN_DIR and VAL_DIR flags depending on the install location of the ImageNet dataset. This is to allow you to train with ImageNet and then evaluate on ImageNet-c.

export DATASET='imagenet'
export TRAIN_DIR=./datasets/imagenet/train
export VAL_DIR=./datasets/imagenet-c/weather/snow/3

Note that when training with ImageNet, the SOURCE flag will not be used.

Architectures, hyperparameters, and optimizers

To select the classifier architecture, you can set the following flags:

export ARCHITECTURE='basic'
export N_CLASSES=10         # number of classes
export SZ=32                # dataset image size (SZ x SZ x 3)
export BS=64                # batch size

The 'basic' architecture is a simple CNN with two convolutional layers and two feed-forward layers. The program will also accept any of the architectures in torchvision.models, including AlexNet and ResNet50. For example, an appropriate ImageNet configuration could be something like this:

export ARCHITECTURE='resnet50'
export N_CLASSES=1000           # number of classes
export SZ=224                   # dataset image size (SZ x SZ x 3)
export BS=64                    # batch size

These flags will also allow you to set the number of output classes for the given architecture, the size of the images in the dataset, and the (training) batch size.

You can also select the optimizer to be used for training the given ARCHITECTURE. Currently, two optimizers are supported: SGD with momentum and AdaDelta. In our paper, we used SGD for the experiments on ImageNet, and AdaDelta for every other experiment (e.g. MNIST, SVHN, CURE-TSR, GTSRB, etc.). To select an optimizer, you can use the --optimizer flag, which currently supports the arguments sgd or adadelta.

Using pre-trained models of natural variation

Yu can set the path to a saved model of natural variation and the dimension of the nuisance space Δ by setting

export MODEL_PATH=./core/models/learned_models/svhn-brightness.pt
export CONFIG_PATH=./core/models/munit/munit.yaml
export DELTA_DIM=8

The CONFIG_PATH should point to a .yaml file with appropriate settings for the MUNIT architecture. An example is given in core/models/munit/munit.yaml. Note that the dimension must match the style_dim parameter in core/models/munit/munit.yaml if you are using the MUNIT framework.

Composing models of natural variation

To compose two models of natural variation, you can simply pass multiple paths after the --model-paths argument. For example, to compose models of contrast and brightness for SVHN, first set

export MODEL_PATH_1=./core/models/learned_models/svhn-brightness.pt
export MODEL_PATH_2=./core/models/learned_models/svhn-contrast.pt

and then add --model-paths $MODEL_PATH_1 $MODEL_PATH_2 to the python command at the bottom of train_bash.sh or train_imagenet.sh.

Training algorithms

By default, the script will train a classifier with the standard ERM formulation. However, by adding flags, you can train classifiers using the three model-based algorithms from our paper (MAT, MRT, and MDA) as well as PGD. For example, to train with MRT and k=10, you can add the flags --mrt -k 10 to the python -m torch.distributed.launch ... command at the bottom of the file. By replacing --mrt with --mat or --mda, you can change the algorithm to MAT or MDA respectively. Similarly, you can use the --pgd flag to train with the PGD algorithm. By default, PGD runs with a step size of 0.01, ɛ=8/255, and 20 steps of gradient ascent.

Distributed settings

You can set the distributed settings with the following flags:

export N_GPUS_PER_NODE=4
export N_NODES=1

This will control how your training is distributed (see torch.nn.parallel.DistributedDataParallel and torch.distributed.launch).

Training models of natural variation

In this work, we used the MUNIT framework to learn models of natural variation. The code that we use to train these models is largely based on the original implementation of MUNIT. To train a model of natural variation with MUNIT, you can run the following shell script:

chmod +x train_munit.sh
./train_munit.sh

You can change the dataset and various directories using the flags in train_munit.sh. In particular, you can set the DATASET and SOURCE environmental variables in the same was as in train_basic.sh and train_imagenet.sh. You can also set various paths, such as the path to the MUNIT configuration file and to directory where you would like to save your output:

export CONFIG_PATH=core/models/munit/munit.yaml
export OUTPUT_PATH=core/models/munit/results

The CONFIG_PATH should point to a .yaml file with appropriate settings for the MUNIT architecture. An example is given in core/models/munit/munit.yaml. Note that the parameter style_dim in this file corresponds to the dimension that will be used for the nuisance space Δ. By default, we have set this to 8, which was the dimension used throughout the experiments section of our paper.

Retrieving a saved model of natural variation

After running train_munit.sh, you can retrieve a saved model of natural variation from ${OUTPUT_PATH}/outputs/munit/checkpoint/gen_<xxxx>.pt, where <xxxx> denotes the iteration number. A module that can be used to reload this .pt file has been provided in the MUNITModelOfNatVar class in core/models/load.py. This class can be used to instantiate a model of natural variation G(x, δ) with a forward pass that takes a batch of images and an appropriately sized nuisance parameter δ. For example, running

G = MUNITModelOfNatVar(args.model_path, reverse=False).cuda()

will return a model of natural variation that can be called in the following way:

imgs, target = next(iter(data_loader))
delta = torch.randn(imgs.size(0), delta_dim, 1, 1).cuda()
mb_images = G(imgs.cuda(), delta)

Here, mb_images will be a batch of images that look semantically similar to imgs but will have different levels of natural variation. Note that delta_dim must be set appropriately in this code snippet to match the style_dim parameter from the .yaml file located at OUTPUT_PATH.

As MUNIT learns mappings in both directions (e.g. from domain A-->B and from domain B--> A), we use the reverse flag to control which direction the MUNIT model maps. By default, reverse=False, meaning that G will map from domain A to B. If reverse=True, G will map from domain B to A.

Using other architectures for models of natural variation

To use other architectures for G, you can simply replace the MUNITModelOfNatVar instantiation in the load_model function in core/models/load.py. In particular, the only requirement is that a model of natural variation should be instantiated as a torch.nn.Module with a forward pass function forward that takes as input a batch of images and a suitably sized nuisance parameter, i.e.

import torch.nn as nn

class MyModel(nn.Module):
  def __init__(self, fname): 
    self.model = self.load_model(fname)

  def forward(x, delta):
    return self.model(x, delta)

  def load_model(fname):
    # Load model from file and return
    return

A library of pre-trained models of natural variation

We provide a library of pre-trained models of natural variation in a public Google Drive folder. In particular, this folder contains models for MNIST, SVHN, GTSRB, CURE-TSR and ImageNet/ImageNet-c:

https://drive.google.com/drive/folders/1J2OQSWuH4VQwWkHDTz_3mM7fQI2Iws6o?usp=sharing

We plan to update this folder frequently with more models of natural variation.

Trouble-shooting

If you run the code in distributed mode over multiple GPUs, you may encounter errors after exiting a program via CTRL-C. Often, this will result in an error that looks something like this:

RuntimeError: Address already in use

If this happens, processes on one or several of your GPUs may still be running. You can check this by running something like

nvidia-smi

If this shows any running processes, you can kill them individually by their PID number. Alternatively, you can kill all running python processes by running

pkill -9 python
Owner
Alex Robey
Alex Robey
PyTorch implementation of Progressive Growing of GANs for Improved Quality, Stability, and Variation.

PyTorch implementation of Progressive Growing of GANs for Improved Quality, Stability, and Variation. Warning: the master branch might collapse. To ob

559 Dec 14, 2022
Predicting Tweet Sentiment Maching Learning and streamlit

Predicting-Tweet-Sentiment-Maching-Learning-and-streamlit (I prefere using Visual Studio Code ) Open the folder in VS Code Run the first cell in requi

1 Nov 20, 2021
Learning to Adapt Structured Output Space for Semantic Segmentation, CVPR 2018 (spotlight)

Learning to Adapt Structured Output Space for Semantic Segmentation Pytorch implementation of our method for adapting semantic segmentation from the s

Yi-Hsuan Tsai 782 Dec 30, 2022
PyTorch implementation of federated learning framework based on the acceleration of global momentum

Federated Learning with Acceleration of Global Momentum PyTorch implementation of federated learning framework based on the acceleration of global mom

0 Dec 23, 2021
Point-NeRF: Point-based Neural Radiance Fields

Point-NeRF: Point-based Neural Radiance Fields Project Sites | Paper | Primary c

Qiangeng Xu 662 Jan 01, 2023
ICNet and PSPNet-50 in Tensorflow for real-time semantic segmentation

Real-Time Semantic Segmentation in TensorFlow Perform pixel-wise semantic segmentation on high-resolution images in real-time with Image Cascade Netwo

Oles Andrienko 219 Nov 21, 2022
ScriptProfilerPy - Module to visualize where your python script is slow

ScriptProfiler helps you track where your code is slow It provides: Code lines t

Lucas BLP 3 Jun 02, 2022
Accelerate Neural Net Training by Progressively Freezing Layers

FreezeOut A simple technique to accelerate neural net training by progressively freezing layers. This repository contains code for the extended abstra

Andy Brock 203 Jun 19, 2022
PyTorch implementation of our ICCV paper DeFRCN: Decoupled Faster R-CNN for Few-Shot Object Detection.

Introduction This repo contains the official PyTorch implementation of our ICCV paper DeFRCN: Decoupled Faster R-CNN for Few-Shot Object Detection. Up

133 Dec 29, 2022
Implement object segmentation on images using HOG algorithm proposed in CVPR 2005

HOG Algorithm Implementation Description HOG (Histograms of Oriented Gradients) Algorithm is an algorithm aiming to realize object segmentation (edge

Leo Hsieh 2 Mar 12, 2022
PyTorch implementation of saliency map-aided GAN for Auto-demosaic+denosing

Saiency Map-aided GAN for RAW2RGB Mapping The PyTorch implementations and guideline for Saiency Map-aided GAN for RAW2RGB Mapping. 1 Implementations B

Yuzhi ZHAO 20 Oct 24, 2022
VisionKG: Vision Knowledge Graph

VisionKG: Vision Knowledge Graph Official Repository of VisionKG by Anh Le-Tuan, Trung-Kien Tran, Manh Nguyen-Duc, Jicheng Yuan, Manfred Hauswirth and

Continuous Query Evaluation over Linked Stream (CQELS) 9 Jun 23, 2022
Multi Camera Calibration

Multi Camera Calibration 'modules/camera_calibration/app/camera_calibration.cpp' is for calculating extrinsic parameter of each individual cameras. 'm

7 Dec 01, 2022
StyleMapGAN - Official PyTorch Implementation

StyleMapGAN - Official PyTorch Implementation StyleMapGAN: Exploiting Spatial Dimensions of Latent in GAN for Real-time Image Editing Hyunsu Kim, Yunj

NAVER AI 425 Dec 23, 2022
Byzantine-robust decentralized learning via self-centered clipping

Byzantine-robust decentralized learning via self-centered clipping In this paper, we study the challenging task of Byzantine-robust decentralized trai

EPFL Machine Learning and Optimization Laboratory 4 Aug 27, 2022
Human Pose estimation with TensorFlow framework

Human Pose Estimation with TensorFlow Here you can find the implementation of the Human Body Pose Estimation algorithm, presented in the DeeperCut and

Eldar Insafutdinov 1.1k Dec 29, 2022
Unrolled Variational Bayesian Algorithm for Image Blind Deconvolution

unfoldedVBA Unrolled Variational Bayesian Algorithm for Image Blind Deconvolution This repository contains the Pytorch implementation of the unrolled

Yunshi HUANG 2 Jul 10, 2022
PyTorch implementation of Decoupling Value and Policy for Generalization in Reinforcement Learning

PyTorch implementation of Decoupling Value and Policy for Generalization in Reinforcement Learning

48 Dec 08, 2022
Real-time Object Detection for Streaming Perception, CVPR 2022

StreamYOLO Real-time Object Detection for Streaming Perception Jinrong Yang, Songtao Liu, Zeming Li, Xiaoping Li, Sun Jian Real-time Object Detection

Jinrong Yang 237 Dec 27, 2022
Physics-informed convolutional-recurrent neural networks for solving spatiotemporal PDEs

PhyCRNet Physics-informed convolutional-recurrent neural networks for solving spatiotemporal PDEs Paper link: [ArXiv] By: Pu Ren, Chengping Rao, Yang

Pu Ren 11 Aug 23, 2022