Auto-Lambda

This repository contains the source code of Auto-Lambda and baselines from the paper, Auto-Lambda: Disentangling Dynamic Task Relationships.

We encourage readers to check out our project page, including more interesting discussions and insights which are not covered in our technical paper.

Multi-task Methods

We implemented all weighting and gradient-based baselines presented in the paper for computer vision tasks: Dense Prediction Tasks (for NYUv2 and CityScapes) and Multi-domain Classification Tasks (for CIFAR-100).

Specifically, we have covered the implementation of these following multi-task optimisation methods:

Weighting-based:

Equal - All task weightings are 1. --weight equal
Uncertainty - https://arxiv.org/abs/1705.07115 --weight uncert
Dynamic Weight Average - https://arxiv.org/abs/1803.10704 --weight dwa
Auto-Lambda - Our approach. --weight autol

Gradient-based:

GradDrop - https://arxiv.org/abs/2010.06808 --grad_method graddrop
PCGrad - https://arxiv.org/abs/2001.06782 --grad_method pcgrad
CAGrad - https://arxiv.org/abs/2110.14048 --grad_method cagrad

Note: Applying a combination of both weighting and gradient-based methods can further improve performance.

Datasets

We applied the same data pre-processing following our previous project: MTAN which experimented on:

NYUv2 [3 Tasks] - 13 Class Segmentation + Depth Estimation + Surface Normal. [288 x 384] Resolution.
CityScapes [3 Tasks] - 19 Class Segmentation + 10 Class Part Segmentation + Disparity (Inverse Depth) Estimation. [256 x 512] Resolution.

Note: We have included a new task: Part Segmentation for CityScapes dataset. Please install the pip install panoptic_parts for CityScapes experiments. The pre-processing file for CityScapes has also been included in the dataset folder.

Experiments

All experiments were written in PyTorch 1.7 and can be trained with different flags (hyper-parameters) when running each training script. We briefly introduce some important flags below.

Flag Name	Usage	Comments
`network`	choose multi-task network: `split, mtan`	both architectures are based on ResNet-50; only available in dense prediction tasks
`dataset`	choose dataset: `nyuv2, cityscapes`	only available in dense prediction tasks
`weight`	choose weighting-based method: `equal, uncert, dwa, autol`	only `autol` will behave differently when set to different primary tasks
`grad_method`	choose gradient-based method: `graddrop, pcgrad, cagrad`	`weight` and `grad_method` can be applied together
`task`	choose primary tasks: `seg, depth, normal` for NYUv2, `seg, part_seg, disp` for CityScapes, `all`: a combination of all standard 3 tasks	only available in dense prediction tasks
`with_noise`	toggle on to add noise prediction task for training (to evaluate robustness in auxiliary learning setting)	only available in dense prediction tasks
`subset_id`	choose domain ID for CIFAR-100, choose `-1` for the multi-task learning setting	only available in CIFAR-100 tasks
`autol_init`	initialisation of Auto-Lambda, default `0.1`	only available when applying Auto-Lambda
`autol_lr`	learning rate of Auto-Lambda, default `1e-4` for NYUv2 and `3e-5` for CityScapes	only available when applying Auto-Lambda

Training Auto-Lambda in Multi-task / Auxiliary Learning Mode:

python trainer_dense.py --dataset [nyuv2, cityscapes] --task [PRIMARY_TASK] --weight autol --gpu 0   # for NYUv2 or CityScapes dataset
python trainer_cifar.py --subset_id [PRIMARY_DOMAIN_ID] --weight autol --gpu 0   # for CIFAR-100 dataset

Training in Single-task Learning Mode:

python trainer_dense_single.py --dataset [nyuv2, cityscapes] --task [PRIMARY_TASK]  --gpu 0   # for NYUv2 or CityScapes dataset
python trainer_cifar_single.py --subset_id [PRIMARY_DOMAIN_ID] --gpu 0   # for CIFAR-100 dataset

Note: All experiments in the original paper were trained from scratch without pre-training.

Benchmark

For standard 3 tasks in NYUv2 (without noise prediction task) in the multi-task learning setting with Split architecture, please follow the results below.

Method	Type	Sem. Seg. (mIOU)	Depth (aErr.)	Normal (mDist.)	Delta MTL
Single	-	43.37	52.24	22.40	-
Equal	W	44.64	43.32	24.48	+3.57%
DWA	W	45.14	43.06	24.17	+4.58%
GradDrop	G	45.39	43.23	24.18	+4.65%
PCGrad	G	45.15	42.38	24.13	+5.09%
Uncertainty	W	45.98	41.26	24.09	+6.50%
CAGrad	G	46.14	41.91	23.52	+7.05%
Auto-Lambda	W	47.17	40.97	23.68	+8.21%
Auto-Lambda + CAGrad	W + G	48.26	39.82	22.81	+11.07%

Note: The results were averaged across three random seeds. You should expect the error range less than +/-1%.

Citation

If you found this code/work to be useful in your own research, please considering citing the following:

@article{liu2022auto_lambda,
    title={Auto-Lambda: Disentangling Dynamic Task Relationships},
    author={Liu, Shikun and James, Stephen and Davison, Andrew J and Johns, Edward},
    journal={Transactions on Machine Learning Research},
    year={2022}
}

Acknowledgement

We would like to thank @Cranial-XIX for his clean implementation for gradient-based optimisation methods.

Contact

If you have any questions, please contact sk.lorenmt@gmail.com.

lorenmt/auto-lambda