Dataset Distillation by Matching Training Trajectories
Project Page | Paper
This repo contains code for training expert trajectories and distilling synthetic data from our Dataset Distillation by Matching Training Trajectories paper (CVPR 2022). Please see our project page for more results.
Dataset Distillation by Matching Training Trajectories George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A. Efros, Jun-Yan Zhu CMU, MIT, UC Berkeley CVPR 2022
The task of “Dataset Distillation” is to learn a small number of synthetic images such that a model trained on this set alone will have similar test performance as a model trained on the full real dataset.
Our method distills the synthetic dataset by directly optimizing the fake images to induce similar network training dynamics as the full, real dataset. We train “student” networks for many iterations on the synthetic data, measure the error in parameter space between the “student” and “expert” networks trained on real data, and back-propagate through all the student network updates to optimize the synthetic pixels.
Wearable ImageNet: Synthesizing Tileable Textures
Instead of treating our synthetic data as individual images, we can instead encourage every random crop (with circular padding) on a larger canvas of pixels to induce a good training trajectory. This results in class-based textures that are continuous around their edges.
Given these tileable textures, we can apply them to areas that require such properties, such as clothing patterns.
Visualizations made using FAB3D
Getting Started
First, download our repo:
git clone https://github.com/GeorgeCazenavette/mtt-distillation.git
cd mtt-distillation
For an express instillation, we include .yaml files.
If you have an RTX 30XX GPU (or newer), run
conda env create -f requirements_11_3.yaml
If you have an RTX 20XX GPU (or older), run
conda env create -f requirements_10_2.yaml
You can then activate your conda environment with
conda activate distillation
Quadro Users Take Note:
torch.nn.DataParallel seems to not work on Quadro A5000 GPUs, and this may extend to other Quadro cards.
If you experience indefinite hanging during training, try running the process with only 1 GPU by prepending CUDA_VISIBLE_DEVICES=0 to the command.
Generating Expert Trajectories
Before doing any distillation, you’ll need to generate some expert trajectories using buffer.py
The following command will train 100 ConvNet models on CIFAR-100 with ZCA whitening for 50 epochs each:
python buffer.py --dataset=CIFAR100 --model=ConvNet --train_epochs=50 --num_experts=100 --zca --buffer_path={path_to_buffer_storage} --data_path={path_to_dataset}
We used 50 epochs with the default learning rate for all of our experts.
Worse (but still interesting) results can be obtained faster through training fewer experts by changing --num_experts. Note that experts need only be trained once and can be re-used for multiple distillation experiments.
Distillation by Matching Training Trajectories
The following command will then use the buffers we just generated to distill CIFAR-100 down to just 1 image per class:
python distill.py --dataset=CIFAR100 --ipc=1 --syn_steps=20 --expert_epochs=3 --max_start_epoch=20 --zca --lr_img=1000 --lr_lr=1e-05 --lr_teacher=0.01 --buffer_path={path_to_buffer_storage} --data_path={path_to_dataset}
ImageNet
Our method can also distill subsets of ImageNet into low-support synthetic sets.
When generating expert trajectories with buffer.py or distilling the dataset with distill.py, you must designate a named subset of ImageNet with the --subset flag.
For example,
python dsitill.py --dataset=ImageNet --subset=imagesquawk --model=ConvNetD5 --ipc=1 --res=128 --syn_steps=20 --expert_epochs=2 --max_start_epoch=10 --lr_img=1000 --lr_lr=1e-06 --lr_teacher=0.01 --buffer_path={path_to_buffer_storage} --data_path={path_to_dataset}
will distill the imagesquawk subset (at 128×128 resolution) into the following 10 images
To register your own ImageNet subset, you can add it to the Config class at the top of utils.py.
Simply create a list with the desired class ID’s and add it to the dictionary.
This gist contains a list of all 1k ImageNet classes and their corresponding numbers.
Acknowledgments
We would like to thank Alexander Li, Assaf Shocher, Gokul Swamy, Kangle Deng, Ruihan Gao, Nupur Kumari, Muyang Li, Gaurav Parmar, Chonghyuk Song, Sheng-Yu Wang, and Bingliang Zhang as well as Simon Lucey’s Vision Group at the University of Adelaide for their valuable feedback. This work is supported, in part, by the NSF Graduate Research Fellowship under Grant No. DGE1745016 and grants from J.P. Morgan Chase, IBM, and SAP. Our code is adapted from https://github.com/VICO-UoE/DatasetCondensation
Reference
If you find our code useful for your research, please cite our paper.
@inproceedings{
cazenavette2022distillation,
title={Dataset Distillation by Matching Training Trajectories},
author={George Cazenavette and Tongzhou Wang and Antonio Torralba and Alexei A. Efros and Jun-Yan Zhu},
booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
year={2022}
}
