Skip to content

Official PyTorch implementation of "DiffusionNAG: Predictor-guided Neural Architecture Generation with Diffusion Models" (ICLR 2024)

Notifications You must be signed in to change notification settings

CownowAn/DiffusionNAG

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

5 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

DiffusionNAG: Predictor-guided Neural Architecture Generation with Diffusion Models

Official Code Repository for the paper DiffusionNAG: Predictor-guided Neural Architecture Generation with Diffusion Models.

Why DiffusionNAG?

  • Existing NAS approaches still result in large waste of time as they need to explore an extensive search space and the property predictors mostly play a passive role such as the evaluators that rank architecture candidates provided by a search strategy to simply filter them out during the search process.

  • We introduce a novel predictor-guided Diffusion-based Neural Architecture Generative framework called DiffusionNAG, which explicitly incorporates the predictors into generating architectures that satisfy the objectives.

  • DiffusionNAG offers several advantages compared with conventional NAS methods, including efficient and effective search, superior utilization of predictors for both NAG and evaluation purposes, and easy adaptability across diverse tasks.

Environment Setup

Create environment with Python 3.7.2 and Pytorch 1.13.1. Use the following command to install the requirements:

cd setup
conda create -n diffusionnag python==3.7.2
conda activate diffusionnag
bash install.sh

Run - NAS-Bench-201

Download datasets, preprocessed datasets, and checkpoints

cd NAS-Bench-201

If you want to run experiments on datasets that are not included in the benchmark, such as "aircraft" or "pets", you will need to download the raw dataset for actually trainining neural architectures on it:

## Download the raw dataset
bash script/download_raw_dataset.sh [DATASET]

We use dataset encoding when training the meta-predictor or performing conditional sampling with the target dataset. To obtain dataset encoding, you need to download the preprocessed dataset:

## Download the preprocessed dataset
bash script/download_preprocessed_dataset.sh

If you want to use the pre-trained score network or meta-predictor, download the checkpoints from the following links.

Download the pre-trained score network and move the checkpoint to checkpoints/scorenet directory:

Download the pre-trained meta-predictor and move the checkpoint to checkpoints/meta_surrogate directory:

Transfer NAG

bash script/transfer_nag.sh [GPU] [DATASET]
## Examples
bash script/transfer_nag.sh 0 cifar10
bash script/transfer_nag.sh 0 cifar100
bash script/transfer_nag.sh 0 aircraft
bash script/transfer_nag.sh 0 pets

Train score network

bash script/tr_scorenet.sh [GPU]

Train meta-predictor

bash script/tr_meta_surrogate.sh [GPU]

Run - MobileNetV3

Download datasets, preprocessed datasets, and checkpoints

cd MobileNetV3

If you want to run experiments on datasets that are not included in the benchmark, such as "aircraft" or "pets", you will need to download the raw dataset for actually trainining neural architectures on it:

## Download the raw dataset
bash script/download_raw_dataset.sh [DATASET]

We use dataset encoding when training the meta-predictor or performing conditional sampling with the target dataset. To obtain dataset encoding, you need to download the preprocessed dataset:

## Download the preprocessed dataset
bash script/download_preprocessed_dataset.sh

If you want to use the pre-trained score network or meta-predictor, download the checkpoints from the following links.

Download the pre-trained score network and move the checkpoint to checkpoints/ofa/score_model directory:

Download the first pre-trained meta-predictor and move the checkpoint to checkpoints/ofa/noise_aware_meta_surrogate directory:

Download the second pre-trained meta-predictor and move the checkpoint to checkpoints/ofa/unnoised_meta_surrogate_from_metad2a directory:

Download the config file for TransferNAG experiments and move the checkpoint to configs directory:

Transfer NAG

bash script/transfer_nag.sh [GPU] [DATASET]
## Examples
bash script/transfer_nag.sh 0,1 cifar10
bash script/transfer_nag.sh 0,1 cifar100
bash script/transfer_nag.sh 0,1 aircraft
bash script/transfer_nag.sh 0,1 pets

Train score network

bash script/tr_scorenet_ofa.sh [GPU]

Train meta-predictor

bash script/tr_meta_surrogate_ofa.sh [GPU]

Citation

If you have found our work helpful for your research, we would appreciate it if you could acknowledge it by citing our work.

@inproceedings{an2023diffusionnag,
  title={DiffusionNAG: Predictor-guided Neural Architecture Generation with Diffusion Models},
  author={An, Sohyun and Lee, Hayeon and Jo, Jaehyeong and Lee, Seanie and Hwang, Sung Ju},
  booktitle={The Twelfth International Conference on Learning Representations},
  year={2023}
}

About

Official PyTorch implementation of "DiffusionNAG: Predictor-guided Neural Architecture Generation with Diffusion Models" (ICLR 2024)

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published