MTAN - 多任务注意力网络

本仓库包含多任务注意力网络（MTAN）及其基准模型的源代码，这些内容出自论文《带有注意力机制的端到端多任务学习》（https://arxiv.org/abs/1803.10704），由刘世坤（Shikun Liu）、爱德华·约翰斯（Edward Johns）和安德鲁·戴维森（Andrew Davison）共同提出。

更多实验结果请参阅我们的项目页面：这里。

最后更新 - 本仓库将不再进行更新。欢迎查看我们的最新工作：Auto-Lambda，其中提供了更多多任务优化方法。

实验

图像到图像预测（一对多）

在 im2im_pred 文件夹下，我们提供了论文中提出的网络架构以及所有基准模型，并基于 NYUv2 数据集进行了实验。所有模型均使用 PyTorch 编写，我们在最新提交中已将实现更新至 PyTorch 1.5 版本。

您可在此下载我们在论文中使用的预处理过的 NYUv2 数据集：这里。我们还使用了来自 这里 的预计算真值法线。原始的 13 类 NYUv2 数据集可以直接从 这个仓库 下载，其分割标签则定义在 这个仓库 中。

很抱歉，由于意外的电脑故障，我无法提供原始的预处理代码。

更新 - 2019年6月：我现在发布了经过预处理的 CityScapes 数据集，包含 2、7 和 19 类语义标签（详情请参阅论文）以及深度（正向和反向）标签。您可以下载 [256x512, 2.42GB] 版本：这里，或 [128x256, 651MB] 版本：这里。

更新 - 2019年10月：对于使用 PyTorch 1.2 的用户：mIoU 评估方法现已更新，以避免在计算二值掩码时出现“零问题”。此外，为正确运行代码，请将 scheduler.step() 移至调用 optimizer.step() 之后，例如在最后一次性能打印步骤之前的一行，以符合更新后的 PyTorch 要求。更多信息请参阅官方 PyTorch 文档 这里。 [此问题已在最新提交中修复。]

更新 - 2020年5月：我们现在提供了官方的 MTAN-DeepLabv3（或类似 ResNet 的架构）设计，以支持更复杂和现代的多任务网络主干。请查看 im2im_pred/model_resnet_mtan 以获取更多详细信息。用户可以轻松地将此模型替换为 im2im_pred 中定义的任何训练模板。

更新 - 2020年7月：我们进一步提升了代码的可读性，并将 im2im_pred 中的所有实现更新至当前最新的 PyTorch 1.5 版本。我们修复了一个错误，用于排除未定义像素的预测，从而更准确地计算语义分割任务中的平均 IoU。此外，我们还为在 NYUv2 数据集中应用数据增强的用户提供了一个额外选项，以避免过拟合并获得更好的性能。

更新 - 2020年11月 [重要！]：我们已更新 mIoU 和像素准确率的公式，使其与 官方 COCO 分割脚本 的标准基准保持一致。现在预计所有方法的 mIoU 性能将提升约 8%。新公式基于所有图像中累积的像素预测来计算 mIoU 和像素准确率，而旧公式则是基于每张图像中像素预测的平均值再取整体平均。

以下表格描述了使用 SegNet（原论文中提出）构建的所有模型（文件）：

文件名	类型	标志	备注
model_segnet_single.py	单任务	task, dataroot	标准单任务学习
model_segnet_stan.py	单任务	task, dataroot	我们的单任务应用方法
model_segnet_split.py	多任务	weight, dataroot, temp, type	多任务学习基准，共享网络在最后一层分离（即硬参数共享）
model_segnet_dense.py	多任务	weight, dataroot, temp	多任务学习基准，每个任务拥有独立的参数空间（即软参数共享）
model_segnet_cross.py	多任务	weight, dataroot, temp	我们的 Cross Stitch Network 实现
model_segnet_mtan.py	多任务	weight, dataroot, temp	我们的多任务学习方法

每个标志的含义如下：

标志名称	用途	备注
task	选择一个任务进行训练：语义（语义分割，深度交叉熵损失）、深度（深度估计，l1 范数损失）或法线（法线预测，余弦相似度损失）	仅适用于单任务学习
dataroot	NYUv2 数据集的目录根路径	放入 im2im_pred 文件夹即可，无需担心 :D
weight	多任务学习的权重选项：等权（直接求和所有任务损失）、DWA（我们的方案）、uncert（我们对 Weight Uncertainty Method 的实现）	仅适用于多任务学习
temp	DWA 权重选项中的超参数温度	用于确定任务权重的软硬程度
type	不同版本的多任务基准拆分：标准、深层、宽层	仅适用于基准拆分
apply_augmentation	打开以在 NYUv2 数据集中应用数据增强，避免过拟合	适用于所有训练模型

要运行任意模型，请先执行 cd im2im_pred/，然后简单地运行 python MODEL_NAME.py --FLAG_NAME 'FLAG_OPTION'（默认选项为不使用数据增强）。若要启用数据增强，请打开 apply_augmentation 标志：python MODEL_NAME.py --FLAG_NAME 'FLAG_OPTION' --apply_augmentation。

请注意，我们在原论文中并未应用任何数据增强。

多任务学习的基准测试

多任务学习的基准测试一直是个棘手的问题，因为每个任务的性能指标和评估方法都不尽相同。在原始论文中，我简单地取了最后10个epoch中各任务性能的平均值，假设我们无法访问验证数据。

为了进行更加标准化和公平的比较，我建议研究人员采用这篇论文第5节、公式4中定义的评估方法，即计算相对于单任务学习的平均相对任务提升。

由于样本量较小，NYUv2数据集很容易出现过拟合。在7月份的更新中，我们提供了一个数据增强选项来缓解过拟合问题（感谢Jialong的帮助）。我们强烈建议在使用该数据增强技术的情况下对NYUv2数据集进行基准测试，以与其他采用相同数据增强技术的SOTA多任务学习方法保持一致，例如PAD-Net和MTI-Net。

视觉十项全能挑战赛（多对多）

我们还提供了针对视觉十项全能挑战赛的源代码，其中我们基于宽残差网络构建了MTAN，并参考了这里的实现。

要运行代码，请按照以下步骤操作：

1. 请从视觉十项全能挑战赛官方网站这里下载数据集和开发工具包。将数据集文件夹decathlon-1.0-data移动到visual_decathlon文件夹下，然后将decathlon_mean_std.pickle移至decathlon-1.0-data文件夹内。

2. 在test文件夹下为每个数据集创建一个子目录，并将所有测试文件移至该子目录中。（这是为了符合PyTorch的数据加载器格式要求。）

3. 在code/coco/PythonAPI文件夹下的十项全能开发工具包中安装setup.py。然后将开发工具包中的pycocotools和annotations文件夹移至visual_decathlon文件夹。

4. 进入visual_decathlon目录，运行python model_wrn_mtan.py --gpu [GPU_ID] --mode [eval, 或 all]进行训练。eval表示在验证集上进行评估（通常用于调试或超参数调优），而all表示在所有数据集上进行训练（通常用于最终评估或基准测试）。

5. 分别运行python model_wrn_eval.py --dataset 'imagenet' 和 'notimagenet'，分别在ImageNet和其他数据集上进行评估。最后，运行python coco_results.py将结果转换为COCO格式，以便在线评估。

其他说明

1. 提供的代码在可读性方面进行了高度优化。如果您发现任何异常行为，请提交issue或直接联系下方的邮箱。
2. 使用提供的代码进行训练，其性能表现（取决于具体任务类型）会与论文中报告的图像到图像预测任务的结果有所不同。不过，模型之间的相对排名保持不变。如果您想比较论文中关于图像到图像预测任务的任何模型，请使用您偏好的训练策略（如学习率、优化器等）重新运行模型，并确保所有训练策略一致，以保证公平性。对于视觉十项全能挑战赛的结果，您可以直接参考论文中给出的数值。若要在您的研究中进行公平比较，请使用相同的主干架构构建您的多任务网络。
3. 根据我的个人经验，在多任务学习中，设计更好的网络架构通常比寻找更优的任务权重分配更为有效且容易实现。

引用

如果您认为本代码/工作对您的研究有所帮助，请考虑引用以下文献：

@inproceedings{liu2019end,
  title={End-to-End Multi-task Learning with Attention},
  author={Liu, Shikun and Johns, Edward and Davison, Andrew J},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
  pages={1871--1880},
  year={2019}
}

致谢

我们感谢Simon Vandenhende在MTAN-DeepLabv3设计方面的帮助；感谢Jialong Wu在MTAN-DeepLabv3基准测试以及NYUv2数据集数据增强实现上的慷慨贡献。

联系方式

如有任何问题，请联系sk.lorenmt@gmail.com。

MTAN (Multi-Task Attention Network) 快速上手指南

本指南基于官方仓库内容整理，旨在帮助开发者快速在 PyTorch 环境下运行 MTAN 模型，进行图像到图像的多任务学习实验。

环境准备

系统要求

• 操作系统: Linux / macOS / Windows
• Python: 3.6+
• PyTorch: 推荐版本 1.5+ (代码已适配 PyTorch 1.5，兼容 1.2+)
• GPU: 推荐使用 NVIDIA GPU 以加速训练

前置依赖

请确保已安装以下核心库：

pip install torch torchvision numpy pillow

若需运行 Visual Decathlon 挑战相关代码，还需安装 pycocotools：

pip install pycocotools

提示：国内用户建议使用清华或阿里镜像源加速安装： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torchvision numpy pillow

安装步骤

本项目无需复杂的安装脚本，直接克隆仓库即可使用。

1. 克隆仓库

   git clone https://github.com/lorenmt/mtan.git
   cd mtan
   

2. 准备数据集 (以 NYUv2 为例)

   代码本身不包含数据集，需手动下载预处理后的数据。

   • 下载地址: NYUv2 Pre-processed Dataset
   • 操作: 下载后解压，将数据文件夹放置在 im2im_pred 目录下，或在运行命令时通过 --dataroot 指定路径。

   注：若需 CityScapes 数据集，可参考 README 中的 Dropbox 链接下载。

基本使用

以下示例展示如何在 NYUv2 数据集上训练标准的 MTAN (SegNet 架构) 多任务模型。

1. 进入实验目录

cd im2im_pred

2. 运行训练命令

使用 model_segnet_mtan.py 启动训练。默认配置为多任务学习，使用 DWA (Dynamic Weight Average) 损失加权策略。

python model_segnet_mtan.py --weight DWA --dataroot '/path/to/your/nyuv2_dataset' --temp 2.0

参数说明：

• --weight: 多任务加权策略。可选值：
  • equal: 损失直接相加
  • DWA: 动态权重平均 (论文提出方法，推荐)
  • uncert: 基于不确定性的加权
• --dataroot: 数据集根目录绝对路径。
• --temp: DWA 策略的温度超参数，控制权重的平滑度 (默认通常为 2.0)。
• --apply_augmentation: (可选) 添加此标志以启用数据增强，防止过拟合 (推荐用于 NYUv2)。

带数据增强的训练示例：

python model_segnet_mtan.py --weight DWA --dataroot '/path/to/your/nyuv2_dataset' --temp 2.0 --apply_augmentation

3. 其他基线模型对比

仓库提供了多种基线模型，只需替换脚本名称即可运行：

• 单任务学习: python model_segnet_single.py --task semantic --dataroot ...
• 硬参数共享 (Hard Parameter Sharing): python model_segnet_split.py --weight equal --dataroot ...
• 软参数共享 (Soft Parameter Sharing): python model_segnet_dense.py --weight equal --dataroot ...
• Cross Stitch Network: python model_segnet_cross.py --weight equal --dataroot ...

4. Visual Decathlon 挑战 (多对多任务)

若需运行 Visual Decathlon 实验，请进入 visual_decathlon 目录：

cd ../visual_decathlon
# 训练所有数据集
python model_wrn_mtan.py --gpu 0 --mode all

注意：运行前需按 README 要求下载 Decathlon 数据集并配置 pycocotools。

---

注意事项：

• 最新提交的代码已修复了 mIoU 计算逻辑，结果应与 COCO 标准基准一致。
• 由于随机性和硬件差异，复现结果可能与论文数值略有波动，但模型间的相对排名应保持一致。