MTL-NAS: Task-Agnostic Neural Architecture Search towards General-Purpose Multi-Task Learning
Official PyTorch Implementation of MTL-NAS
Please refer to our paper for more technical details:
Yuan Gao*, Haoping Bai*, Zequn Jie, Jiayi Ma, Kui Jia, Wei Liu. MTL-NAS: Task-Agnostic Neural Architecture Search towards General-Purpose Multi-Task Learning, IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020. [arXiv]
If this code is helpful to your research, please consider citing our paper by:
@inproceedings{mtlnas2020,
title={MTL-NAS: Task-Agnostic Neural Architecture Search towards General-Purpose Multi-Task Learning},
author={Yuan Gao and Haoping Bai and Zequn Jie and Jiayi Ma and Kui Jia and Wei Liu},
year={2020},
booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}
}
Setup
Install the necessary dependencies:
$ pip install -r requirements.txt
Dataset
Follow the instruction here to prepare the dataset. Alternatively, download the preprocessed dataset here.
Download the converted PyTorch models from here, then create a weights
directory and unzip the models inside.
When you are all set, you should have the following file structure:
datasets/nyu_v2/list
datasets/nyu_v2/nyu_v2_mean.npy
datasets/nyu_v2/nyu_train_val
weights/vgg_deeplab_lfov/tf_deeplab.pth
weights/nyu_v2/tf_finetune_seg.pth
weights/nyu_v2/tf_finetune_normal.pth
Training
All the arguments to train/eval MTLNAS are shown in core/config/defaults.py
. The configuration files for different experiments are also provided in the configs
directory. To run the NDDR-CNN baseline with VGG-16 architecture, simply call:
$ CUDA_VISIBLE_DEVICES=0 python tools/train.py --config-file configs/vgg/vgg_nyuv2_nddr.yaml
To run MTLNAS training with default configuration, call:
$ CUDA_VISIBLE_DEVICES=0 python tools/train_nas.py --config-file configs/ablation/vgg_nyuv2_default.yaml
Evaluation
To evaluate the final checkpoint for the NDDR-CNN baseline experiment, call:
$ CUDA_VISIBLE_DEVICES=0 python tools/eval.py --config-file configs/vgg/vgg_nyuv2_nddr.yaml
To evaluate the final checkpoint for default MTLNAS, call:
$ CUDA_VISIBLE_DEVICES=0 python tools/eval_nas.py --config-file configs/ablation/vgg_nyuv2_default.yaml
You can download and extract the final checkpoint for default MTLNAS to ckpts
directory and evaluate it by running the command above.