Cross-Task Affinity Learning for Multitask Dense Scene Predictions

This is the official PyTorch implementation of the paper: "Cross-Task Affinity Learning for Multitask Dense Scene Predictions" (WACV 2025).

CTAL enhances task refinement in multitask learning (MTL) by capturing both local and long-range cross-task interactions, improving performance for CNN and Transformer-based architectures.

To clone the repo and setup the environment, run:

git clone https://github.com/Armanfard-Lab/CTAL.git
cd CTAL
conda env create -f environment.yml
conda activate ctal

Before running experiments, download (and extract) the datasets:

1. NYUv2
2. Cityscapes
3. PASCAL_Context

For CNN-based models:

1. Download weights HRNet-W18-C-Small-v2.
2. Save them to code_cnn/models/pretrained_models/hrnet_w18_small_model_v2.pth.

For Transformer-based models, the pretrained swin_v2 backbones are dowloaded automatically with torchvision.

We provide two implementations:

• code_cnn/ → CNN-based code (supports Cityscapes, NYUv2)
• code_transformer/ → Transformer-based code (supports PASCAL_Context)

1. Enter the appropriate directory:

   cd code_cnn

   or

   cd code_transformer

2. Set the db_root (dataroot) in configs/mypath.py to where you stored the dataset.
3. Update --storage_root and --config in train.sh
4. Run:

   bash train.sh

If you use this code or find our work helpful, please cite:

@INPROCEEDINGS{10943529,
  author={Sinodinos, Dimitrios and Armanfard, Narges},
  booktitle={2025 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)}, 
  title={Cross-Task Affinity Learning for Multitask Dense Scene Predictions}, 
  year={2025},
  volume={},
  number={},
  pages={1546-1555},
  keywords={Computer vision;Refining;Computer architecture;Predictive models;Transformers;Multitasking;multitask learning;computer vision;dense scene predictions},
  doi={10.1109/WACV61041.2025.00158}
}

Multitask learning (MTL) has become prominent for its ability to predict multiple tasks jointly, achieving better pertask performance with fewer parameters than single-task learning. Recently, decoder-focused architectures have significantly improved multitask performance by refining task predictions using features from related tasks. However, most refinement methods struggle to efficiently capture both local and long-range dependencies between task-specific representations and cross-task patterns. In this paper, we introduce the Cross-Task Affinity Learning (CTAL) module, a lightweight framework that enhances task refinement in multitask networks. CTAL effectively captures local and long-range cross-task interactions by optimizing task affinity matrices for parameter-efficient grouped convolutions without concern for information loss. Our results demonstrate state-of-the-art MTL performance for both CNN and transformer backbones, using significantly fewer parameters than single-task learning.

This repo borrows several elements from Multi-Task-Learning-PyTorch and MTAN.