多任务学习

本仓库旨在使用 PyTorch 实现多种多任务学习模型及训练策略。代码库补充了以下工作：

密集预测任务中的多任务学习：综述

Simon Vandenhende、Stamatios Georgoulis、Wouter Van Gansbeke、Marc Proesmans、Dengxin Dai 和 Luc Van Gool。

MTI-Net：用于多任务学习的多尺度任务交互网络

Simon Vandenhende、Stamatios Georgoulis 和 Luc Van Gool。

关于多任务学习的最新文献列表可以在此处找到：链接。

研讨会

:loudspeaker: :loudspeaker: :loudspeaker: 我们在 ICCV 2021 上组织了一场关于多任务学习的 研讨会（链接）。

• 1月13日：我们的特邀报告录像现已在 Youtube 上发布。

安装

代码可在较新的 PyTorch 版本上运行，例如 1.4。假设已安装 Anaconda，最重要的包可以通过以下命令安装：

conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
conda install imageio scikit-image		   	   # 图像处理
conda install -c conda-forge opencv		           # OpenCV
conda install pyyaml easydict                 		   # 配置文件
conda install termcolor                       		   # 彩色打印语句

我们参考 requirements.txt 文件来了解我们自身环境中各包的具体版本。

使用说明

设置

为了在您自己的机器上运行代码，需要修改以下文件：

• 修改 utils/mypath.py 中的数据集路径，例如 /path/to/pascal/。
• 在 configs/your_env.yml 中指定输出目录。所有结果都将存储在此目录下。
• 进行边缘评估需要 seism 仓库。请参阅 ./evaluation/seism/ 中的 README 文件。
• 如果您想使用 HRNet 主干网络，请从 这里 下载预训练权重。 提供的配置文件使用 HRNet-18 主干。请下载 hrnet_w18_small_model_v2.pth 并将其保存到 ./models/pretrained_models/ 目录中。

首次运行代码时，数据集将自动下载到指定路径。

训练

用于训练模型的配置文件位于 configs/ 目录中。可通过以下命令进行训练：

python main.py --config_env configs/env.yml --config_exp configs/$DATASET/$MODEL.yml

评估

我们在训练结束时评估最佳模型。评估标准基于我们综述论文中的公式 10，需要事先预训练一组单任务网络。为加快训练速度，可以在配置文件中添加以下行，仅在最后 10 个 epoch 进行评估：

eval_final_10_epochs_only: True

支持

以下数据集和任务得到支持。

数据集	语义分割	深度估计	法线估计	边缘检测	显著性检测	人体部位
PASCAL	Y	N	Y	Y	Y	Y
NYUD	Y	Y	辅助	辅助	N	N

以下模型得到支持。

主干网络	HRNet	ResNet
单任务	Y	Y
多任务	Y	Y
Cross-Stitch		Y
NDDR-CNN		Y
MTAN		Y
PAD-Net	Y
MTI-Net	Y

参考文献

本代码库大量借鉴了 ASTMT 仓库。尤其是评估部分和数据加载器均来自该仓库。

引用

如果您觉得本仓库对您的研究有所帮助，请考虑引用以下文献：

@article{
  author={S. Vandenhende and S. Georgoulis and W. Van Gansbeke and M. Proesmans and D. Dai and L. Van Gool},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 
  title={Multi-Task Learning for Dense Prediction Tasks: A Survey}, 
  year={2021},
  volume={},
  number={},
  pages={1-1},
  doi={10.1109/TPAMI.2021.3054719}}

@article{vandenhende2020mti,
  title={MTI-Net: Multi-Scale Task Interaction Networks for Multi-Task Learning},
  author={Vandenhende, Simon and Georgoulis, Stamatios and Van Gool, Luc},
  journal={ECCV2020},
  year={2020}
}

@InProceedings{MRK19,
  Author    = {Kevis-Kokitsi Maninis and Ilija Radosavovic and Iasonas Kokkinos},
  Title     = {Attentive Single-Tasking of Multiple Tasks},
  Booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  Year      = {2019}
}

@article{pont2015supervised,
  title={Supervised evaluation of image segmentation and object proposal techniques},
  author={Pont-Tuset, Jordi and Marques, Ferran},
  journal={IEEE transactions on pattern analysis and machine intelligence},
  year={2015},
}

更新

更多信息请参见 issue #1。

初始代码使用了 ASTMT 中的 NYUDv2 数据加载器。该实现与我们在综述中进行实验时所用的不同。因此，我们重新编写了 NYUDv2 数据加载器，以确保与综述结果一致。为避免问题，最好删除旧版本的 NYUDv2 数据集。之后，当您再次使用 NYUDv2 数据集时，Python 脚本将自动下载正确的版本。

深度估计任务采用逐像素方式评估，以与综述保持一致。这与 ASTMT 的做法不同，后者是将结果在图像间取平均值。

许可证

本软件采用知识共享许可协议发布，仅允许个人和科研用途。如需商业授权，请联系作者。许可证摘要可在此查看：链接。

致谢

作者感谢丰田公司通过 TRACE 项目以及 MACCHINA（鲁汶大学，C14/18/065）提供的支持。

Multi-Task-Learning-PyTorch 快速上手指南

本指南旨在帮助开发者快速部署并运行基于 PyTorch 的多任务学习（Multi-Task Learning）模型库。该库实现了多种多任务学习策略及模型架构（如 MTI-Net, PAD-Net 等），支持语义分割、深度估计、法向量预测等密集预测任务。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.6+
• GPU: 支持 CUDA 的 NVIDIA 显卡（可选，但强烈推荐用于训练）
• 包管理器: 推荐使用 Anaconda 进行环境管理

前置依赖说明： 本项目核心依赖 PyTorch。此外，还需要 opencv、imageio、scikit-image 等图像处理库，以及 pyyaml、easydict 用于配置解析。

国内加速建议： 安装 Python 包时，建议使用清华或阿里镜像源以加快下载速度。 例如：pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <package_name>

安装步骤

1. 创建并激活 Conda 环境

首先创建一个独立的虚拟环境，并安装 PyTorch 及相关基础依赖。以下命令基于 CUDA 10.2 版本（可根据实际显卡驱动调整）：

conda create -n mtl_pytorch python=3.8
conda activate mtl_pytorch

# 安装 PyTorch 全家桶 (使用官方源或国内镜像)
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch

# 安装图像处理与工具库
conda install imageio scikit-image
conda install -c conda-forge opencv
conda install pyyaml easydict
conda install termcolor

注：若需精确复现原作者环境，可参考项目根目录下的 requirements.txt。

2. 克隆代码仓库

git clone https://github.com/SimonVandenhende/Multi-Task-Learning-PyTorch.git
cd Multi-Task-Learning-PyTorch

3. 配置路径与预训练模型

在运行代码前，需修改部分配置文件以适应本地环境：

1. 数据集路径：编辑 utils/mypath.py，将数据集路径修改为您本地的存储位置（例如：/path/to/pascal/）。
   • 提示：首次运行时，代码会自动下载数据集到指定路径。
2. 输出目录：编辑 configs/your_env.yml，指定结果保存目录。
3. HRNet 骨干网络（可选）：若配置文件中使用了 HRNet backbone，需手动下载预训练权重。
   • 下载地址：HRNet-Image-Classification
   • 获取 hrnet_w18_small_model_v2.pth 文件。
   • 将其保存至项目目录：./models/pretrained_models/。
4. 边缘评估工具（可选）：如需进行边缘检测评估，需克隆 seism 仓库，具体参见 ./evaluation/seism/README。

基本使用

训练模型

配置文件位于 configs/ 目录下。您可以根据需要选择数据集和模型架构。

执行以下命令启动训练（以 PASCAL 数据集和某模型为例）：

python main.py --config_env configs/env.yml --config_exp configs/$DATASET/$MODEL.yml

参数说明：

• --config_env: 指向包含输出路径等全局设置的环境配置文件。
• --config_exp: 指向具体的实验配置（包含数据集、模型结构、超参数等）。

加速评估技巧

默认情况下，评估基于论文中的公式并在训练结束后对最佳模型进行。为了节省时间，您可以在配置文件中添加以下行，仅在最后 10 个 epoch 进行评估：

eval_final_10_epochs_only: True

支持的任务与模型

当前版本支持以下组合：

• 数据集: PASCAL (支持语义分割、法向量、边缘、显著性、人体部件), NYUD (支持语义分割、深度)。
• 骨干网络: HRNet, ResNet。
• 多任务架构: Single-Task, Multi-Task (硬共享), Cross-Stitch, NDDR-CNN, MTAN, PAD-Net, MTI-Net。

运行完成后，结果将保存在您在 configs/your_env.yml 中指定的输出目录下。