LibMTL

LibMTL 是一个基于 PyTorch 构建的开源多任务学习（MTL）库。有关详细介绍和 API 使用说明，请参阅 最新文档。

:star: 在 GitHub 上给我们点个赞吧——这对我们来说是很大的鼓励！

:bangbang: 一篇关于 基于梯度的多目标深度学习 的全面综述现已在 arXiv 上发布，同时还附带了一个 精彩列表。快来查看吧！

新闻

• [2025年4月21日] 增加了对 UPGrad 的支持。
• [2025年2月18日] 增加了对双层优化方法 Auto-Lambda（TMLR 2022）的支持。
• [2025年2月17日] 增加了对 FAMO（NeurIPS 2023）、SDMGrad（NeurIPS 2023）以及 MoDo（NeurIPS 2023；JMLR 2024）的支持。
• [2025年2月6日] 增加了对两种双层优化方法的支持：MOML（NeurIPS 2021；AIJ 2024）、FORUM（ECAI 2024）。
• [2024年9月19日] 增加了对 FairGrad（ICML 2024）的支持。
• [2024年8月31日] 增加了对 ExcessMTL（ICML 2024）的支持。
• [2024年7月24日] 增加了对 STCH（ICML 2024）的支持。
• [2024年2月8日] 增加了对 DB-MTL 的支持。
• [2023年8月16日]：增加了对 MoCo（ICLR 2023）的支持。非常感谢作者 @heshandevaka 的帮助。
• [2023年7月11日] 论文被 JMLR 接受。
• [2023年6月19日] 增加了对 Aligned-MTL（CVPR 2023）的支持。
• [2023年3月10日]：新增了 QM9 和 PAWS-X 示例。
• [2022年7月22日]：增加了对 Nash-MTL（ICML 2022）的支持。
• [2022年7月21日]：增加了对 Learning to Branch（ICML 2020）的支持。非常感谢 @yuezhixiong 的贡献（#14）。
• [2022年3月29日]：论文现已在 arXiv 上发布。

目录

• 功能特性
• 整体框架
• 支持的算法
• 支持的基准数据集
• 安装指南
• 快速入门
  • 下载数据集
  • 运行模型
• 引用
• 贡献者
• 联系我们
• 致谢
• 许可证

功能特性

• 统一性：LibMTL 提供了一个统一的代码库来实现多种多任务学习算法，并采用一致的评估流程，涵盖数据处理、指标计算和超参数设置等环节。该库支持多个具有代表性的多任务学习基准数据集，从而能够对不同算法进行定量、公平且一致的比较。
• 全面性：LibMTL 支持众多最先进的多任务学习方法，包括8种架构和16种优化策略。同时，它还提供了对多个覆盖不同领域的基准数据集的公平比较。
• 可扩展性：LibMTL 遵循模块化设计原则，允许用户灵活便捷地添加自定义组件或进行个性化修改。因此，借助 LibMTL 的支持，用户可以轻松快速地开发新的优化策略和架构，或将现有多任务学习算法应用于新的应用场景。

整体框架

各模块的详细介绍请参见 文档。

支持的算法

目前，LibMTL 支持以下算法：

优化策略	会议/来源	参数选项
等权法 (EW)	-	--weighting EW
梯度归一化 (GradNorm)	ICML 2018	--weighting GradNorm
不确定性权重法 (UW)	CVPR 2018	--weighting UW
MGDA (官方代码)	NeurIPS 2018	--weighting MGDA
动态权重平均法 (DWA) (官方代码)	CVPR 2019	--weighting DWA
几何损失策略 (GLS)	CVPR 2019研讨会	--weighting GLS
冲突梯度投影法 (PCGrad)	NeurIPS 2020	--weighting PCGrad
梯度符号丢弃法 (GradDrop)	NeurIPS 2020	--weighting GradDrop
公平多任务学习法 (IMTL)	ICLR 2021	--weighting IMTL
梯度疫苗法 (GradVac)	ICLR 2021	--weighting GradVac
冲突规避梯度下降法 (CAGrad）(官方代码)	NeurIPS 2021	--weighting CAGrad
MOML	NeurIPS 2021	--weighting MOML
Nash-MTL（官方代码）	ICML 2022	--weighting Nash_MTL
随机损失加权法 (RLW)	TMLR 2022	--weighting RLW
Auto-Lambda （官方代码）	TMLR 2022	--weighting AutoLambda
MoCo	ICLR 2023	--weighting MoCo
Aligned-MTL（官方代码）	CVPR 2023	--weighting Aligned_MTL
FAMO （官方代码）	NeurIPS 2023	--weighting FAMO
SDMGrad （官方代码）	NeurIPS 2023	--weighting SDMGrad
MoDo （官方代码）	NeurIPS 2023	--weighting MoDo
FORUM	ECAI 2024	--weighting FORUM
STCH （官方代码）	ICML 2024	--weighting STCH
ExcessMTL （官方代码）	ICML 2024	--weighting ExcessMTL
FairGrad （官方代码）	ICML 2024	--weighting FairGrad
DB-MTL	arXiv	--weighting DB_MTL
UPGrad（官方代码）	arXiv	--weighting UPGrad

架构	会议	参数
硬参数共享（HPS）	ICML 1993	--arch HPS
十字绣网络（Cross_stitch）	CVPR 2016	--arch Cross_stitch
多门混合专家模型（MMoE）	KDD 2018	--arch MMoE
多任务注意力网络（MTAN）（官方代码）	CVPR 2019	--arch MTAN
自定义门控机制（CGC）、渐进式分层提取（PLE）	ACM RecSys 2020	--arch CGC, --arch PLE
学习分支（LTB）	ICML 2020	--arch LTB
DSelect-k（官方代码）	NeurIPS 2021	--arch DSelect_k

支持的基准数据集

数据集	问题	任务数量	任务	多输入	支持的主干网络
NYUv2	场景理解	3	语义分割+ 深度估计+ 表面法线预测	✘	ResNet50/ SegNet
Cityscapes	场景理解	2	语义分割+ 深度估计	✘	ResNet50
Office-31	图像识别	3	分类	✓	ResNet18
Office-Home	图像识别	4	分类	✓	ResNet18
QM9	分子性质预测	11（默认）	回归	✘	GNN
PAWS-X	同义句识别	4（默认）	分类	✓	Bert

安装

1. 创建虚拟环境

   conda create -n libmtl python=3.10
   conda activate libmtl
   

2. 克隆仓库

   git clone https://github.com/median-research-group/LibMTL.git
   

3. 安装 LibMTL

   pip install torch==2.3.0+cu121 torchvision==0.18.0+cu121 -f https://download.pytorch.org/whl/torch_stable.html
   pip install torch-scatter==2.1.2+pt23cu121 torch_sparse==0.6.18+pt23cu121 -f https://pytorch-geometric.com/whl/torch-2.3.0+cu121.html
   
   cd LibMTL
   pip install -r requirements.txt
   pip install -e .
   

快速入门

我们以 NYUv2 数据集为例，展示如何使用 LibMTL。

下载数据集

我们使用的 NYUv2 数据集由 mtan 预处理。您可以从 这里下载该数据集。

运行模型

NYUv2 数据集的完整训练代码位于 examples/nyu 中。文件 main.py 是在 NYUv2 数据集上进行训练的主文件。

您可以通过运行以下命令查看命令行参数。

python main.py -h

例如，运行以下命令将在 NYUv2 数据集上训练一个使用 EW 和 HPS 的多任务学习模型。

python main.py --weighting EW --arch HPS --dataset_path /path/to/nyuv2 --gpu_id 0 --scheduler step --mode train --save_path PATH

更多详细信息请参阅 文档。

引用

如果您发现 LibMTL 对您的研究或开发有所帮助，请引用以下内容：

@article{lin2023libmtl,
  title={{LibMTL}: A {P}ython Library for Multi-Task Learning},
  author={Baijiong Lin and Yu Zhang},
  journal={Journal of Machine Learning Research},
  volume={24},
  number={209},
  pages={1--7},
  year={2023}
}

贡献者

LibMTL 由 Baijiong Lin 开发并维护。

联系我们

如果您有任何问题或建议，请随时通过 提交 issue 或发送邮件至 bj.lin.email@gmail.com 与我们联系。

致谢

我们谨向以下公开仓库的作者表示感谢（按字母顺序排列）：CAGrad、dselect_k_moe、MultiObjectiveOptimization、mtan、MTL、nash-mtl、pytorch_geometric 和 xtreme。

许可证

LibMTL 采用 MIT 许可证发布。

LibMTL 快速上手指南

LibMTL 是一个基于 PyTorch 构建的开源多任务学习（Multi-Task Learning, MTL）库。它提供了统一的代码库和评估流程，支持多种主流 MTL 算法和数据集，旨在帮助开发者公平、一致地比较和复现不同的多任务学习策略。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS, 或 Windows
• Python 版本：3.8 - 3.11 (推荐 3.9+)
• 核心依赖：
  • PyTorch (版本需与您的 CUDA 环境匹配)
  • numpy, scipy, pandas 等基础科学计算库

提示：建议先安装好 PyTorch，LibMTL 会自动检测并适配已安装的 PyTorch 版本。

安装步骤

您可以通过 PyPI 直接安装稳定版，或从 GitHub 安装最新开发版。国内用户推荐使用清华或阿里镜像源以加速下载。

方式一：通过 PyPI 安装（推荐）

# 使用默认源
pip install LibMTL

# 或使用国内镜像源（推荐中国开发者）
pip install LibMTL -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装（获取最新功能）

git clone https://github.com/median-research-group/LibMTL.git
cd LibMTL
pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

LibMTL 的设计遵循模块化原则，最简单的使用方式是通过命令行运行内置的示例脚本。以下以经典的 NYUv2 数据集和 GradNorm 算法为例。

1. 下载数据集

LibMTL 提供了自动下载脚本。首次运行时，数据将被下载到本地缓存目录。

python examples/nyu/main.py --data_path ./data/nyu --download

2. 运行模型

使用 --weighting 参数指定多任务优化策略。例如，运行 GradNorm 算法：

python examples/nyu/main.py \
    --data_path ./data/nyu \
    --weighting GradNorm \
    --gpu_id 0 \
    --seed 0

参数说明：

• --data_path: 数据集存放路径。
• --weighting: 选择优化策略（如 EW, GradNorm, PCGrad, Nash_MTL 等，完整列表见官方文档）。
• --gpu_id: 指定使用的 GPU 编号。
• --seed: 随机种子，确保结果可复现。

3. 自定义开发

若需修改网络结构或添加新算法，LibMTL 允许您通过继承基类轻松扩展。以下是一个极简的代码调用示例：

from LibMTL import Trainer
from LibMTL.weighting import GradNorm

# 配置参数
args = {
    'weighting': 'GradNorm',
    'data_path': './data/nyu',
    'gpu_id': 0,
    # ... 其他参数
}

# 初始化并训练
trainer = Trainer(args)
trainer.train()

更多高级用法、支持的算法列表及 API 详情，请参阅 LibMTL 官方文档。