OmniNet：多模态多任务学习的统一架构

OmniNet 是一种用于多模态多任务学习的统一且扩展版 Transformer 架构。单一的 OmniNet 架构可以编码来自几乎任何现实世界领域的多种输入（文本、图像、视频），并能够在广泛的任务范围内进行异步多任务学习。OmniNet 架构由多个称为神经外设的子网络组成，这些子网络用于将特定领域的输入编码为时空表示，并连接到一个名为中央神经处理器（CNP）的通用中央神经网络。CNP 实现了一个基于 Transformer 的通用时空编码器和一个多任务解码器。

本仓库包含 OmniNet：多模态多任务学习的统一架构（Pramanik 等人）的官方 PyTorch 实现。论文展示了一个 OmniNet 实例，该实例经过联合训练，能够同时执行词性标注、图像字幕生成、视觉问答和视频动作识别等任务。我们还开源了预训练模型，以便复现论文中展示的结果。

设置

要求

• 最低硬件：8GB 内存 + NVIDIA GPU（8GB 以上）
• 基于 Linux 的操作系统
• NVIDIA 驱动程序 410+
• Anaconda 包管理器
• 其余依赖将由 Anaconda 自动安装（见下文）

安装

1. 克隆代码库。
2. 使用提供的文件创建并激活环境。（这将安装所有必需的依赖项，包括 CUDA 和 cuDNN）

$ conda env create -f environment.yml
$ source activate omninet

下载数据集

我们提供了一个下载脚本，用于下载并预处理论文中提到的所有任务（除 PENN 词性标注外）的训练和测试数据。下载的数据将存储在 /data 文件夹中。

$ python scripts/init_setup.py

预训练模型

论文中展示的各种单任务和多任务模型的预训练模型可以从以下 URL 获取，可用于评估架构性能或对外部数据进行预测。

1. VQA：https://storage.googleapis.com/omninet/200k_vqa.zip
2. 字幕生成：https://storage.googleapis.com/omninet/280k_caption.zip
3. HMDB：https://storage.googleapis.com/omninet/11k_hmdb.zip
4. PENN：https://storage.googleapis.com/omninet/70k_penn.zip
5. VQA+HMDB+字幕生成+PENN：https://storage.googleapis.com/omninet/all_15k.zip

训练

训练脚本包含用于在 COCO 图像字幕生成、VQA、HMDB 动作识别以及 PENN 词性标注任务上训练模型的子程序。首先，请按照上述说明使用 scripts/init_setup.py 下载所有所需数据集。训练脚本既可以用于单任务训练，也可以用于多个不同任务的异步训练。它实现了多 GPU 的 hogwild 训练模式，允许在多个 GPU 上同时训练单个模型以完成多项任务。

$ python train.py <训练迭代次数> <任务列表> <批大小> --n_gpus <GPU 数量> \
--save_interval <保存间隔> --eval_interval <评估间隔>

支持的任务列表包括 vqa、hmdb、caption 和 penn。（然而，init_setup.py 中的数据下载工具不会下载 PENN 标注的数据集，因此需要手动下载并预处理 PENN 词性标注数据集，才能训练模型进行词性标注。）

例如，要在视觉问答这一单一任务上训练模型，进行 100,000 次迭代，批大小为 128：

$ python train.py 100000 vqa 128 --n_gpus 1 --save_interval 500 --eval_interval 500

若要异步训练多个任务，只需用逗号分隔任务列表和对应的批大小即可。例如，在三个 GPU 上联合训练 VQA、HMDB 和 COCO 字幕生成任务，批大小分别为 128、64 和 128：

$ python train.py 100000 vqa,hmdb,caption 128,64,128 --n_gpus 3 --save_interval 500 --eval_interval 500

当同时训练多个任务时，建议使用多 GPU 系统；如果 GPU 显存不足，则可适当降低每个任务的批大小。

评估

评估脚本可用于评估任何给定任务的预训练模型。要评估从上述预训练模型链接下载的任意模型文件，需先解压下载的模型压缩包，然后运行以下脚本：

$ python evaluate.py <任务名称 > < model.pth 文件路径 > --batch_size < 批大小 >

支持的任务名称如下：

• vqa：计算 VQA 验证集上的得分
• hmdb：计算 HMDB 第一版训练/测试分割上的准确率
• caption：计算 COCO 验证集上的得分

预测

预测脚本可用于对论文中提及的所有任务的外部数据进行预测，也可用于对未训练过的任务进行零样本预测。支持的任务包括 vqa、hmdb、caption 和 penn。当前的预测脚本仅实现贪心解码，未来我们将引入束搜索以提升性能。请从上述 URL 下载并解压多任务预训练模型。

$ wget https://storage.googleapis.com/omninet/all_15k.zip
$ unzip all_15k.zip

对预训练任务的预测：

词性标注

$ python predict.py model.pth penn --text "there is no dark side of the moon really, as a matter of fact its all dark"
词性标注预测：EX VBZ RB VBN IN DT NN IN NN PRP$ DT NN IN NN PRP$ DT JJ

图像字幕生成

$ python predict.py model.pth caption --image sample_data/kids-hand-signals.jpg 
字幕生成预测：一个人正在城市街道上骑自行车

视觉问答

$ python predict.py model.pth vqa --image sample_data/kids-hand-signals.jpg --text "孩子在驾驶什么？"
视觉问答预测：自行车

视频动作识别

$ python predict.py model.pth hmdb --video sample_data/riding_horse.avi 
动作识别预测：骑马

未训练任务上的预测（迈向零样本机器学习）：

由于中央神经处理器采用了共享的多模态表征学习架构，OmniNet 也可以用于对其从未训练过的任务进行零样本预测。例如，尽管该模型从未在这些任务上进行过训练，但其多模态架构仍可用于视频字幕生成和视频问答。以下结果展示了对一段骑马视频进行视频字幕生成和视频问答的示例。（注：零样本学习的结果目前仍处于实验阶段，并非总是具有实际意义）

视频字幕生成

同样的预测脚本也可用于视频字幕生成。不过，此时我们不再提供图像，而是直接提供原始视频文件进行预测：

$ python predict.py model.pth caption --video sample_data/riding_horse.avi
字幕预测：一个人在靠近水域的海滩上骑马。

视频问答

$ python predict.py model.pth vqa --video sample_data/riding_horse.avi --text "他在做什么？"
视频问答预测：骑马

$ python predict.py model.pth vqa --video sample_data/riding_horse.avi --text "里面是男人还是女人？"
视频问答预测：男人

$ python predict.py model.pth vqa --video sample_data/riding_horse.avi --text "这匹马是什么颜色的？"
视频问答预测：棕色

作者

1. Subhojeet Pramanik（IBM Watson）email@subho.in
2. Priyanka Agrawal（IBM Research）pagrawal.ml@gmail.com
3. Aman Hussain（阿姆斯特丹大学）email@amanhussain.com
4. Sayan Dutta duttasayan.7@gmail.com

引用

我们欢迎研究界进一步拓展我们的工作。如果您觉得这段代码有用，请引用我们的论文：OmniNet：一种用于多模态多任务学习的统一架构

@ARTICLE{2019arXiv190707804P,
       author = {{Pramanik}, Subhojeet and {Agrawal}, Priyanka and {Hussain}, Aman},
        title = "{OmniNet: 一种用于多模态多任务学习的统一架构}",
      journal = {arXiv e-prints},
     keywords = {计算机科学——机器学习、计算机科学——计算与语言、计算机科学——计算机视觉与模式识别、计算机科学——神经与进化计算、统计学——机器学习},
         year = "2019",
        month = "7月",
          eid = {arXiv:1907.07804},
        pages = {arXiv:1907.07804},
archivePrefix = {arXiv},
       eprint = {1907.07804},
 primaryClass = {cs.LG},
       adsurl = {https://ui.adsabs.harvard.edu/abs/2019arXiv190707804P},
      adsnote = {由SAO/NASA天体物理数据系统提供}
}

OmniNet 快速上手指南

OmniNet 是一个用于多模态多任务学习的统一架构，基于 Transformer 扩展而成。单个模型即可同时处理文本、图像和视频输入，并支持异步多任务训练（如词性标注、图像描述、视觉问答和视频动作识别）。

环境准备

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

• 操作系统: Linux
• 内存: 至少 8GB RAM
• GPU: NVIDIA GPU (显存 8GB+)
• 驱动: NVIDIA Driver 410+
• 包管理器: Anaconda

注意：本项目依赖 CUDA 和 cuDNN，后续安装步骤将通过 Conda 自动配置。

安装步骤

1. 克隆仓库

git clone <repo_url>
cd <repo_directory>

(请将 <repo_url> 和 <repo_directory> 替换为实际的项目地址和目录名)

2. 创建并激活环境

使用提供的 environment.yml 文件创建隔离环境，该操作会自动安装 PyTorch、CUDA 及其他依赖项：

conda env create -f environment.yml
source activate omninet

3. 下载数据集

运行初始化脚本以下载并预处理论文中提及的大部分任务数据（除 PENN 词性标注外），数据将保存至 /data 目录：

python scripts/init_setup.py

注：若需训练 PENN 词性标注任务，需手动下载并预处理该数据集。

基本使用

以下展示如何使用预训练模型进行最简单的**预测（Prediction）**操作。

1. 下载预训练模型

下载包含所有任务的多任务预训练模型并解压：

wget https://storage.googleapis.com/omninet/all_15k.zip
unzip all_15k.zip

2. 运行预测示例

OmniNet 支持对未训练过的任务进行零样本（Zero-shot）预测。以下是几个典型用例：

图像描述生成 (Image Captioning)

python predict.py model.pth caption --image sample_data/kids-hand-signals.jpg 

输出示例：a person riding a bike on a city street

视觉问答 (Visual Question Answering)

python predict.py model.pth vqa --image sample_data/kids-hand-signals.jpg --text "What is the kid driving?"

输出示例：bike

视频动作识别 (Video Action Recognition)

python predict.py model.pth hmdb --video sample_data/riding_horse.avi 

输出示例：ride_horse

零样本视频描述 (Zero-shot Video Captioning) 即使模型未在视频描述任务上专门训练，也可尝试直接预测：

python predict.py model.pth caption --video sample_data/riding_horse.avi

输出示例：a person riding a horse on a beach near a body of water.

---

如需进行训练或评估，请参考官方文档中的 train.py 和 evaluate.py 使用说明。