This repository contains the PyTorch implementation of the paper SNAS: stochastic neural architecture search, ICLR 2019.
By Xie Sirui, Zheng Hehui, Liu Chunxiao, Lin Liang.
Python >= 3.5.5, PyTorch == 0.4, torchvision == 0.2.0
CIFAR-10 can be automatically downloaded by torchvision, ImageNet needs to be downloaded manually.
Single GPU setup to search the architecture with 8 cells:
python train_search.py --snas --epochs 150 --seed 6 --layer 8 --init_channels 16 --temp 1 \
--temp_min 0.03 --nsample 1 --temp_annealing --resource_efficient \
--resource_lambda 1.5e-3 --log_penalty --drop_path_prob 3e-1 --method 'reparametrization' \
--loss --remark "snas_order_layer_8_batch_64_drop_0.3_error_lnR_1e-2_reparam_gpu_1" &
Multi GPU setup to search the architecture with 20 cells:
python -m torch.distributed.launch --nproc_per_node=8 train_search.py --snas --epochs 150 --seed 6 --layer 20 --init_channels 36 --batch_size 72 --temp 1 --temp_min 0.03 \
--nsample 1 --distributed --temp_annealing --resource_efficient --resource_lambda 1.5e-3 --log_penalty \
--drop_path_prob 2e-1 --method 'reparametrization' --loss --remark "snas_order_layer_20_batch_14_drop_0.1_error_lnR_1.5e-3_reparam_gpu_8" &
Note that we implement the resource constraint with three methods, i.e., reparametrization, discrete and policy gradient.
python train_edge_all.py --auxiliary --cutout --arch {arch_edge_all}
Note that {arch_edge_all} could be SNAS_mild_edge_all, SNAS_moderate_edge_all, SNAS_aggressive_edge_all in genotypes.py.
Tensorboard visualization:
tensorboard --logdir=runs/ --port={port_num}
Note that all the experiments above will save the tensorboard log file in runs/ directory
| Model | Top-1 | Download | MD5 |
|---|---|---|---|
| cifar-10 mild constraint | 97.02% | [Google Drive] [Baidu Pan (pin:an86)] | aebc72d18b2c79925fafcebcd68ec745 |
Evaluate the trained model:
python test_edge_all.py --auxiliary --arch {arch} --model_path {model_path}
If you find our codes or trined models useful in your research, please consider to star our repo and cite our paper:
@inproceedings{
xie2018snas,
title={{SNAS}: stochastic neural architecture search},
author={Sirui Xie and Hehui Zheng and Chunxiao Liu and Liang Lin},
booktitle={International Conference on Learning Representations},
year={2019},
url={https://openreview.net/forum?id=rylqooRqK7},
}