Single Path One-Shot

Single Path One-Shot by Megvii Research.

Introduction

This repository provides the implementation of Single Path One-Shot Neural Architecture Search with Uniform Sampling.

Our Trained Model / Checkpoint

• OneDrive: Link

Supernet

Our trained Supernet weight is in $Link/Supernet/checkpoint-150000.pth.tar, which can be used by Search.

Search

Our search result is in $Link/Search/checkpoint.pth.tar, which can be used by Evaluation.

Evaluation

Out searched models have been trained from scratch, is can be found in $Link/Evaluation/$ARCHITECTURE.

Here is a summary:

Usage

1. Setup Dataset and Flops Table

Download the ImageNet Dataset and move validation images to labeled subfolders. To do this, you can use the following script: https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh

Download the flops table to accelerate Flops calculation which is required in Uniform Sampling. It can be found in $Link/op_flops_dict.pkl.

We recommend to create a folder data and use it in both Supernet training and Evaluation training.

Here is a example structure of data:

data
|--- train                 ImageNet Training Dataset
|--- val                   ImageNet Validation Dataset
|--- op_flops_dict.pkl     Flops Table

2. Train Supernet

Train supernet with the following command:

cd src/Supernet
python3 train.py --train-dir $YOUR_TRAINDATASET_PATH --val-dir $YOUR_VALDATASET_PATH

3. Search in Supernet with Evolutionary Algorithm

Search in supernet with the following command:

cd src/Search
python3 search.py

It will use ../Supernet/checkpoint-latest.pth.tar as Supernet's weight, please make sure it exists or modify the path manually.

4. Get Searched Architecture

Get searched architecture with the following command:

cd src/Evaluation
python3 eval.py

It will generate folder in data/$YOUR_ARCHITECTURE. You can train the searched architecture from scratch in the folder.

5. Train from Scratch

Finally, train and evaluate the searched architecture with the following command.

Train:

cd src/Evaluation/data/$YOUR_ARCHITECTURE
python3 train.py --train-dir $YOUR_TRAINDATASET_PATH --val-dir $YOUR_VALDATASET_PATH

Evaluate:

cd src/Evaluation/data/$YOUR_ARCHITECTURE
python3 train.py --eval --eval-resume $YOUR_WEIGHT_PATH --train-dir $YOUR_TRAINDATASET_PATH --val-dir $YOUR_VALDATASET_PATH

Citation

If you use these models in your research, please cite:

@article{guo2019single,
        title={Single path one-shot neural architecture search with uniform sampling},
        author={Guo, Zichao and Zhang, Xiangyu and Mu, Haoyuan and Heng, Wen and Liu, Zechun and Wei, Yichen and Sun, Jian},
        journal={arXiv preprint arXiv:1904.00420},
        year={2019}
}