DensePred

Learning Universal Representations for Multiple Dense Prediction Tasks.

Dependencies

• Python 3.6+
• PyTorch 1.0 (or newer version)
• torchvision 0.2.2 (or newer version)
• progress
• matplotlib
• numpy

Usage

Prepare dataset

We use the preprocessed NYUv2 dataset provided by MTAN. Download the dataset and place the dataset folder in ./data/

Our method

Train Single Task Learning Network with SegNet:

CUDA_VISIBLE_DEVICES=<gpu-id> python nyu_stl.py --method single-task --backbone segnet --head segnet_head --task <task, e.g. semantic, depth, normal> --out ./results/stl --dataroot ./data/nyuv2

Train the Universal Representation Network:

CUDA_VISIBLE_DEVICES=<gpu-id> python nyu_url.py --backbone segnet --head segnet_head --tasks semantic depth normal --out ./results/url --single-dir ./results/stl --dataroot ./data/nyuv2

Baselines

We also provide code nyu_mtl_baseline.py for several multi-task learning optimization strategies, including uniformly weighting [uniform]; MGDA [mgda], adapted from the source code; GradNorm [gradnorm]; DWA [dwa], from the source code, Uncertainty Weighting [uncert]; Gradient Surgery [gs] from Pytorch-PCGrad; IMTL [imtl_l, imtl_g, imtl_h]; GradDrop [gd]; and CAGrad [ca] from the source code.

Train a Multi-task Learning Network using MGDA with SegNet:

CUDA_VISIBLE_DEVICES=<gpu-id> python nyu_mtl_baseline.py --backbone segnet --head segnet_head --tasks semantic depth normal --out ./results/mtl-baselines --dataroot ./data/nyuv2 --weight #weight: uniform, mgda, gradnorm, dwa, uncert, gs, imtl_l, imtl_g, imtl_h, gd, ca

Model Zoo

STL Models | URL Model

Acknowledge

We thank authors of Multi-Task-Learning-PyTorch, MTAN, MGDA, Pytorch-PCGrad, CAGrad for their code!

Contact

For any question, you can contact Wei-Hong Li.

Citation

If you use this code, please cite our papers:

@article{li2023Universal,
    author    = {Li, Wei-Hong and Liu, Xialei and Bilen, Hakan},
    title     = {Universal Representations: A Unified Look at Multiple Task and Domain Learning},
    journal   = {International Journal of Computer Vision},
    pages     = {1--25},
    year      = {2023},
    publisher = {Springer}
}

@inproceedings{li2021Universal,
    author    = {Li, Wei-Hong and Liu, Xialei and Bilen, Hakan},
    title     = {Universal Representation Learning From Multiple Domains for Few-Shot Classification},
    booktitle = {IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2021},
    pages     = {9526-9535}
}

@inproceedings{li2020knowledge,
    author    = {Li, Wei-Hong and Bilen, Hakan},
    title     = {Knowledge distillation for multi-task learning},
    booktitle = {European Conference on Computer Vision (ECCV) Workshop},
    year      = {2020},
}