This repository provides the official PyTorch implementation of UNAS that was presented at CVPR 2020. The paper presents results in two search spaces including DARTS and ProxylessNAS spaces. Our paper can be found here.

UNAS is examined using Python 3.5. To install the requirements, follow these steps:

1. Run

$ pip install -r requirements.txt

2. Install apex using:

$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --no-cache-dir ./

We examined UNAS on the CIFAR-10, CIFAR-100, and ImageNet datasets. The CIFAR datasets will be downloaded automatically if you don't have a local copy. However, you may need to download the ImageNet 2012 dataset and structure the dataset under train and val subfloders. You can follow this page to structure the dataset. The data directory should be in the form:

data/
    ├── train/
        ├── n01440764/
        ├── n01443537/
        ├── ...
    ├── val/
        ├── n01440764/
        ├── n01443537/
        ├── ...        

In UNAS, in the search phase, we first search for either normal/reduction cells (in DARTS space) or the whole network architecture (in ProxylessNAS space). Then in the evaluation phase, in order to measure the performance of search, we train the discovered cells or network architectures from scratch. The commands for performing both search and evaluation are provided in commands.sh. Note that we use 4 GPUs for search on CIFAR-10 and CIFAR-100, 1 GPU for evaluation on these datasets, and 8 GPUs for search and eval on ImageNet.

In order to run the evaluation code with the cells discovered in our experiments, you can set --genotype to UNAS_CIFAR10, UNAS_CIFAR100, or UNAS_IMAGENET when evaluating in DARTS space. Similarly, you can set --genotype to UNAS when evaluating in ProxylessNAS space (check the example command in commands.sh.

You can monitor search and evaluation progress using tensorboar. You can launch tensorboard using these commands:

$ export EXPR_ID=UNIQUE_EXPR_ID
$ export CHECKPOINT_DIR=PATH_TO_CHECKPOINT_DIR
$ tensorboard --logdir $CHECKPOINT_DIR/search-$EXPR_ID/       # During search
$ tensorboard --logdir $CHECKPOINT_DIR/eval-$EXPR_ID/         # During evaluation

Please check the LICENSE file.

@inproceedings{vahdat2019UNAS,
  title={{UNAS}: Differentiable Architecture Search Meets Reinforcement Learning},
  author={Vahdat, Arash and Mallya, Arun and Liu, Ming-Yu and  Kautz, Jan},
  booktitle = {Computer Vision and Pattern Recognition (CVPR)},
  year={2020}
}