Unlabeled Data Improves Adversarial Robustness
This repository contains code for reproducing data and models from the NeurIPS 2019 paper Unlabeled Data Improves Adversarial Robustness by Yair Carmon, Aditi Raghunathan, Ludwig Schmidt, Percy Liang and John C. Duchi.
CIFAR-10 unlabeled data and trained models
Below are links to files containing our unlabeled data and pretrained models:
- 500K unlabeled data from TinyImages (with pseudo-labels)
- Trained heuristic defense RST_adv(50K+500K) model (see Table 1 in the paper)
- Trained certified defense RST_stab(50K+500K) model (see Figure 1 in the paper)
Additional files:
- Trained data sourcing model
(Classifies between CIFAR-10 and non-CIFAR-10 content) - TinyImages indices with keywords matching CIFAR-10
(from the CIFAR-10.1 paper) - Nearest neighbor L2 distances between CIFAR-10 test set and TinyImages
Dependencies
To create a conda environment called semisup-adv containing all the dependencies, run
conda env create -f environment.yml
Note: We tested this code on 2 GPUs in parallel, each with 12GB of memory. Running on CPUs or GPUs with less memory might require adjustments.
The code in this repo is based on code from the following sources:
- TRADES: https://github.com/yaodongyu/TRADES
- Randomized smoothing: https://github.com/locuslab/smoothing
- AutoAugment: https://github.com/DeepVoltaire/AutoAugment
- Cutout: https://github.com/uoguelph-mlrg/Cutout/blob/master/util/cutout.py
- FoolBox: https://github.com/bethgelab/foolbox
- ShakeShake: https://github.com/hysts/pytorch_shake_shake
- CIFAR-10.1: https://github.com/modestyachts/CIFAR-10.1
Running robust self-training
To run robust self-training you will a pickle file containing pseudo-labeled data. You can download ti_500K_pseudo_labeled.pickle
containing our 500K pseudo-labeled TinyImages, or you can generate one from scratch using the instructions below.
Adversarial training with TRADES
The following command performs adversarial training and produces a model
equivalent to RST_adv(50K+500K) described in the paper.
python robust_self_training.py --aux_data_filename ti_500K_pseudo_labeled.pickle --distance l_inf --epsilon 0.031 --model_dir rst_adv
When the script finishes running there will a be checkpoint file called rst_adv/checkpoint-epoch200.pt
. The following commands runs a PGD attack (PGD_Ours from the paper) on the model
python attack_evaluation.py --model_path rst_adv/checkpoint-epoch200.pt --attack pgd --output_suffix pgd_ours
To run the Carlini-Wanger attack on randomly selected 1000 images from the test set, use
python attack_evaluation.py --model_path rst_adv/checkpoint-epoch200.pt --attack cw --output_suffix cw --num_eval_batches 5 --shuffle_testset
Stability training
The following commands performs stability training and produces a model equivalent to
RST_stab(50K+500K) described in the paper.
python robust_self_training.py --aux_data_filename ti_500K_pseudo_labeled.pickle --distance l_2 --epsilon 0.25 --model_dir rst_stab --epochs 800
When the script finishes running there will a be checkpoint file called rst_stab/checkpoint-epoch800.pt
. The following commands runs randomized smoothing certification on the model, as described in the paper.
python smoothing_evaluation.py --model_path rst_stab/checkpoint-epoch800.pt --sigma 0.25
Creating the unlabeled data from scratch
Note: creating the unlabeled data from scratch takes a while; plan for three days at least.
Step zero: Downloading data
Create a data directory that has the following files:
Step one: Tiny Image preliminaries
In this step, we do the following two preliminary steps.
- Compute distances from all the TinyImages to CIFAR-10 test set, in order to ensure we do not add any images from the test set to the unlabeled data sourced from TinyImages.
- Create train/test data for selection model (See Appendix B.6)
Note that the data directory should contain the following files: tiny_images.bin
, cifar10_keywords_unique_v7.json
, tinyimage_subset_indices_v7.json
and tinyimage_subset_data_v7.pickle
.
Here is an example run.
python tinyimages_preliminaries.py --data_dir ../data/ --output_dir ../data
Step two: Train a selection model
Here we train the data selection model described in Appendix B.6 of the paper. Note that data_dir
should contain the following files: tiny_images.bin
, ti_vs_cifar_inds.pickle
(from above).
Here is an example run.
python train_cifar10_vs_ti.py --output_dir ../cifar10-vs-ti/ --data_dir ../data/
Step three: Selecting unlabeled data and removing CIFAR-10 test set
We apply the model trained above on TinyImages and select images based on the predictions, while making sure to remove all images that are close (in l2 distance) to the CIFAR-10 test set.
python tinyimages_prediction.py --model_path ../cifar10-vs-ti/model_state_epoch520.pth --data_dir ../data --output_dir ../data/ --output_filename ti_500K_unlabeled.pickle
Step four: Training a vanilla model on CIFAR-10
We now train a model (Wide ResNet 28-10) on CIFAR-10 training set.
python robust_self_training.py --distance l_2 --beta 0 --unsup_fraction 0 --model_dir vanilla
Step five: Generating pseudo-labels
As a final step, we generate pseudo-labels by applying the classifier from Step 4 on the unlabeled data sourced in Step 3.
python generate_pseudolabels.py --model_dir ../vanilla --model_epoch 200 --data_dir ../data/ --data_filename ti_500K_unlabeled.pickle --output_dir ../data/ --output_filename ti_500K_pseudo_labeled.pickle
Reference
@inproceedings{carmon2019unlabeled,
author = {Yair Carmon and Aditi Raghunathan and Ludwig Schmidt and Percy Liang and John Duchi},
title = {Unlabeled Data Improves Adversarial Robustness},
year = 2019,
booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
}