UNITER: 通用图像-文本表示学习

这是 UNITER（ECCV 2020）的官方仓库。 目前该仓库支持在以下数据集上对 UNITER 进行微调： NLVR2、VQA、VCR、 SNLI-VE、针对 COCO 和 Flickr30k 的图像-文本检索任务， 以及 Referring Expression Comprehensions（RefCOCO、RefCOCO+ 和 RefCOCO-g）。 我们同时发布了 UNITER-base 和 UNITER-large 的预训练检查点。 此外，还提供了使用领域内数据进行预训练的 UNITER-base 模型。

本仓库中的部分代码源自并修改自以下开源实现： PyTorch、 HuggingFace、 OpenNMT 和 Nvidia。 图像特征则使用 BUTD 提取。

环境要求

为了便于复现，我们提供了 Docker 镜像。请安装以下软件：

• NVIDIA 驱动程序（418 及以上版本）；
• Docker（19.03 及以上版本）；
• NVIDIA 容器工具包。

我们的脚本要求用户属于 docker 用户组，以便无需使用 sudo 即可运行 Docker 命令。 我们仅支持配备 NVIDIA GPU 的 Linux 系统。我们已在 Ubuntu 18.04 和 V100 显卡上进行了测试。 由于采用了混合精度训练，建议使用具备 Tensor Core 的 GPU。

快速入门

注意：请运行 bash scripts/download_pretrained.sh $PATH_TO_STORAGE 来获取我们最新的预训练检查点。这将下载基础模型和大型模型。

我们以 NLVR2 作为使用本代码库的端到端示例。

1. 使用以下命令下载处理好的数据和预训练模型：

   bash scripts/download_nlvr2.sh $PATH_TO_STORAGE
   

   下载完成后，您应看到如下文件夹结构：

   ├── ann
   │   ├── dev.json
   │   └── test1.json
   ├── finetune
   │   ├── nlvr-base
   │   └── nlvr-base.tar
   ├── img_db
   │   ├── nlvr2_dev
   │   ├── nlvr2_dev.tar
   │   ├── nlvr2_test
   │   ├── nlvr2_test.tar
   │   ├── nlvr2_train
   │   └── nlvr2_train.tar
   ├── pretrained
   │   └── uniter-base.pt
   └── txt_db
       ├── nlvr2_dev.db
       ├── nlvr2_dev.db.tar
       ├── nlvr2_test1.db
       ├── nlvr2_test1.db.tar
       ├── nlvr2_train.db
       └── nlvr2_train.db.tar
   

2. 启动用于运行实验的 Docker 容器：

   # Docker 镜像应自动拉取
   source launch_container.sh $PATH_TO_STORAGE/txt_db $PATH_TO_STORAGE/img_db \
       $PATH_TO_STORAGE/finetune $PATH_TO_STORAGE/pretrained
   

   启动脚本会尊重 $CUDA_VISIBLE_DEVICES 环境变量。请注意，源代码被挂载到容器内的 /src 目录下，而不是直接构建到镜像中，这样用户修改后无需重新构建镜像即可生效。（数据文件夹则单独挂载到容器中，以提高文件夹结构的灵活性。）

3. 对 NLVR2 任务进行微调：

   # 在容器内
   python train_nlvr2.py --config config/train-nlvr2-base-1gpu.json
   
   # 更多自定义选项
   horovodrun -np $N_GPU python train_nlvr2.py --config $YOUR_CONFIG_JSON
   

4. 对 NLVR2 任务进行推理并评估：

   # 推理
   python inf_nlvr2.py --txt_db /txt/nlvr2_test1.db/ --img_db /img/nlvr2_test/ \
       --train_dir /storage/nlvr-base/ --ckpt 6500 --output_dir . --fp16
   
   # 评估
   # 在 Docker 外部运行此命令（经测试适用于 Python 3.6）
   # 或将标注 JSON 文件复制到挂载的文件夹中
   python scripts/eval_nlvr2.py ./results.csv $PATH_TO_STORAGE/ann/test1.json
   

   上述命令将对我们训练的模型进行推理。您可以根据需要替换 --train_dir 和 --ckpt 参数，使用第 3 步中训练的自有模型。目前我们仅支持单 GPU 推理。

5. 自定义设置：

   # 训练选项
   python train_nlvr2.py --help
   

   • 命令行参数会覆盖 JSON 配置文件；
   • JSON 配置会覆盖 argparse 的默认值。
   • 使用 horovodrun 运行多 GPU 训练；
   • --gradient_accumulation_steps 可模拟多 GPU 训练。

6. 其他操作：

   # 文本标注预处理
   bash scripts/create_txtdb.sh $PATH_TO_STORAGE/txt_db $PATH_TO_STORAGE/ann
   
   # 图像特征提取（经测试可在 Titan-Xp 上运行；可能无法在最新 GPU 上运行）
   bash scripts/extract_imgfeat.sh $PATH_TO_IMG_FOLDER $PATH_TO_IMG_NPY
   
   # 图像预处理
   bash scripts/create_imgdb.sh $PATH_TO_IMG_NPY $PATH_TO_STORAGE/img_db
   

   如果您希望完整复现整个预处理流程，可以参考这些脚本。

下游任务微调

VQA

注意：训练和推理应在 Docker 容器内进行。

1. 下载数据：

   bash scripts/download_vqa.sh $PATH_TO_STORAGE
   

2. 训练：

   horovodrun -np 4 python train_vqa.py --config config/train-vqa-base-4gpu.json \
       --output_dir $VQA_EXP
   

3. 推理：

   python inf_vqa.py --txt_db /txt/vqa_test.db --img_db /img/coco_test2015 \
       --output_dir $VQA_EXP --checkpoint 6000 --pin_mem --fp16
   

   结果文件将保存至 $VQA_EXP/results_test/results_6000_all.json，可用于提交至评估服务器。

VCR

注意：训练和推理应在 Docker 容器内进行。

1. 下载数据：

   bash scripts/download_vcr.sh $PATH_TO_STORAGE
   

2. 训练：

   horovodrun -np 4 python train_vcr.py --config config/train-vcr-base-4gpu.json \
       --output_dir $VCR_EXP
   

3. 推理：

   horovodrun -np 4 python inf_vcr.py --txt_db /txt/vcr_test.db \
       --img_db "/img/vcr_gt_test/;/img/vcr_test/" \
       --split test --output_dir $VCR_EXP --checkpoint 8000 \
       --pin_mem --fp16
   

   结果文件将保存至 $VCR_EXP/results_test/results_8000_all.csv，可用于提交至 VCR 排行榜进行评估。

VCR 第二阶段预训练

注意：预训练应在 Docker 容器内运行

1. 如果尚未下载 VCR 数据，请先下载：

   bash scripts/download_vcr.sh $PATH_TO_STORAGE
   

2. 第二阶段预训练：

   horovodrun -np 4 python pretrain_vcr.py --config config/pretrain-vcr-base-4gpu.json \
       --output_dir $PRETRAIN_VCR_EXP
   

视觉蕴含（SNLI-VE）

注意：训练应在 Docker 容器内运行

1. 下载数据：

   bash scripts/download_ve.sh $PATH_TO_STORAGE
   

2. 训练：

   horovodrun -np 2 python train_ve.py --config config/train-ve-base-2gpu.json \
       --output_dir $VE_EXP
   

图像-文本检索

下载数据：

bash scripts/download_itm.sh $PATH_TO_STORAGE

注意：图像-文本检索计算量较大，尤其是在 COCO 数据集上。

零样本图像-文本检索（Flickr30k）

# 每个图像-文本对都需要进行排序；请尽可能使用多张 GPU
horovodrun -np $NGPU python inf_itm.py \
    --txt_db /txt/itm_flickr30k_test.db --img_db /img/flickr30k \
    --checkpoint /pretrain/uniter-base.pt --model_config /src/config/uniter-base.json \
    --output_dir $ZS_ITM_RESULT --fp16 --pin_mem

图像-文本检索（Flickr30k）

• 普通微调：

  horovodrun -np 8 python train_itm.py --config config/train-itm-flickr-base-8gpu.json
  

• 使用难负样本的微调：

  horovodrun -np 16 python train_itm_hard_negatives.py \
      --config config/train-itm-flickr-base-16gpu-hn.jgon
  

图像-文本检索（COCO）

• 使用难负样本的微调：

  horovodrun -np 16 python train_itm_hard_negatives.py \
      --config config/train-itm-coco-base-16gpu-hn.json
  

指代表达

1. 下载数据：

   bash scripts/download_re.sh $PATH_TO_STORAGE
   

2. 训练：

   python train_re.py --config config/train-refcoco-base-1gpu.json \
       --output_dir $RE_EXP
   

3. 推理与评估：

   source scripts/eval_refcoco.sh $RE_EXP
   

   结果文件将被写入 $RE_EXP/results_test/ 目录下。

同样地，运行 RefCOCO+/RefCOCOg 时需更改相应的配置文件和脚本。

预训练

下载：

bash scripts/download_indomain.sh $PATH_TO_STORAGE

预训练：

horovodrun -np 8 python pretrain.py --config config/pretrain-indomain-base-8gpu.json \
    --output_dir $PRETRAIN_EXP

遗憾的是，由于 CC/SBU 特征数据体积过大，我们无法托管这些数据。用户需要自行处理这些数据。我们将在近期提供一个较小的样本，以便于参考预期的格式。

引用

如果您在研究中使用了本代码，请考虑引用以下文献：

@inproceedings{chen2020uniter,
  title={Uniter: Universal image-text representation learning},
  author={Chen, Yen-Chun and Li, Linjie and Yu, Licheng and Kholy, Ahmed El and Ahmed, Faisal and Gan, Zhe and Cheng, Yu and Liu, Jingjing},
  booktitle={ECCV},
  year={2020}
}

许可证

MIT

UNITER 快速上手指南

UNITER 是一个通用的图像 - 文本表示学习模型，支持视觉问答 (VQA)、图像 - 文本检索、视觉蕴含等多种多模态任务。本指南基于官方仓库整理，帮助开发者快速在 Linux 环境下部署并运行示例。

环境准备

本项目依赖 Docker 容器化运行，仅支持 Linux + NVIDIA GPU 环境（推荐 Ubuntu 18.04 + V100/A100 等具备 Tensor Cores 的显卡以启用混合精度训练）。

请确保宿主机已安装以下组件：

• NVIDIA Driver: 版本 418+
• Docker: 版本 19.03+
• NVIDIA Container Toolkit: 用于在 Docker 中调用 GPU

此外，当前用户需加入 docker 用户组，以便无需 sudo 执行 docker 命令：

sudo usermod -aG docker $USER
# 执行后需重新登录或重启终端生效

安装步骤

1. 获取预训练模型与示例数据

以 NLVR2 任务为例，首先下载处理好的数据及预训练权重。请将 $PATH_TO_STORAGE 替换为你本地的存储路径。

# 下载最新预训练检查点 (包含 base 和 large 模型)
bash scripts/download_pretrained.sh $PATH_TO_STORAGE

# 下载 NLVR2 任务所需的数据和微调配置
bash scripts/download_nlvr2.sh $PATH_TO_STORAGE

下载完成后，目录结构应如下所示：

$PATH_TO_STORAGE
├── ann             # 标注文件
├── finetune        # 微调配置
├── img_db          # 图像特征数据库
├── pretrained      # 预训练模型 (.pt)
└── txt_db          # 文本数据库

2. 启动 Docker 容器

使用官方提供的脚本启动容器，该脚本会自动拉取镜像并将代码和数据挂载到容器中。

source launch_container.sh $PATH_TO_STORAGE/txt_db $PATH_TO_STORAGE/img_db \
    $PATH_TO_STORAGE/finetune $PATH_TO_STORAGE/pretrained

注意：脚本会尊重 $CUDA_VISIBLE_DEVICES 环境变量来指定使用的 GPU。

基本使用

以下演示如何在容器内完成 NLVR2 任务的微调训练与推理评估。

1. 开始微调训练

单卡训练示例：

# 进入容器后执行
python train_nlvr2.py --config config/train-nlvr2-base-1gpu.json

多卡训练示例（需安装 Horovod）：

horovodrun -np $N_GPU python train_nlvr2.py --config $YOUR_CONFIG_JSON

2. 运行推理与评估

使用训练好的模型（或默认提供的检查点）进行测试集推理：

# 推理阶段 (在容器内运行)
python inf_nlvr2.py --txt_db /txt/nlvr2_test1.db/ --img_db /img/nlvr2_test/ \
    --train_dir /storage/nlvr-base/ --ckpt 6500 --output_dir . --fp16

评估阶段需在容器外运行（需 Python 3.6+），或将标注文件复制到挂载目录：

# 评估阶段 (在宿主机运行)
python scripts/eval_nlvr2.py ./results.csv $PATH_TO_STORAGE/ann/test1.json

3. 其他任务

对于 VQA、VCR、图像 - 文本检索等其他任务，流程类似：

1. 运行对应的 scripts/download_*.sh 下载数据。
2. 在容器内运行 train_*.py 进行微调。
3. 运行 inf_*.py 进行推理。

具体命令可参考项目根目录下的 Downstream Tasks Finetuning 章节。