场景图生成网络——RelTR(TPAMI2023)

一、ReITR概述

场景图是一种图结构,其节点代表图像中的实体,边代表实体间的关系。这项技术超越了传统的对象检测,与视觉关系检测紧密相关,对图像检索、图像字幕、视觉问答(VQA)和图像生成等多种视觉-语言任务具有潜在的应用价值。

不同于传统的两阶段方法(先检测对象再分类关系),RelTR 采用一种一阶段方法,直接基于视觉外观预测场景图,无需合并实体或标记所有可能的谓词。这种方法基于DETR(Detection Transformer)的启发,将对象检测视为一种端到端的集合预测任务,并提出了一种基于二分匹配的集合预测损失函数。

RelTR 的主要贡献如下:

  • 相比大多数现有方法使用的对象检测骨架分类密集关系,RelTR 能够通过解码视觉外观及从数据中学习的主体和客体查询来生成稀疏场景图。

  • RelTR 仅基于视觉外观生成场景图,具有较少的参数和更快的推理速度,同时达到了最先进的性能。

  • 设计了一种集合预测损失来执行真实三元组和预测三元组之间的匹配,采用基于IoU的分配策略。

  • 通过解耦实体注意力,RelTR 的三元组解码器能够改进主体和客体的定位和分类。

论文还对在 Visual Genome、Open Images V6 和 VRD 数据集上的实验结果进行了详细讨论,证明了模型的优越性能和快速推理能力。RelTR 的源代码和预训练模型已公开提供,以便研究人员和开发者使用和进一步开发。

二、ReITR安装使用

  1. 安装

下载 RelTR 仓库:

python 复制代码
git clone https://github.com/yrcong/RelTR.git
cd RelTR

如果您想对一张图片进行推断,只需要 python=3.6, PyTorch=1.6 和 matplotlib!您可以按照以下方式配置环境:

python 复制代码
# 创建一个 conda 环境
conda create -n reltr python=3.6
conda activate reltr

# 安装包
conda install pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.1 -c pytorch
conda install matplotlib

如果您想在 Visual Genome 上训练/评估 RelTR,您需要更多的准备:

a) 需要 Scipy(我们使用的是 1.5.2)和 pycocotools。

python 复制代码
conda install scipy
pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

b) 按照 data 目录中的 README 准备数据集。

c) 一些常用的评估代码(IoU)需要编译... 我们将用 Pytorch 代码替换它。

python 复制代码
# 编译计算盒交集的代码
cd lib/fpn
sh make.sh

目录结构如下:

python 复制代码
RelTR
| 
│
└───data
│   └───vg
│       │   rel.json
│       │   test.json
│       │   train.json
|       │   val.json
|       │   images
│   └───oi
│       │   rel.json
│       │   test.json
│       │   train.json
|       │   val.json
|       │   images
└───datasets    
... 
  1. 使用方法

推理

a) 下载我们在 Visual Genome 数据集上预训练的 RelTR 模型,并将其放在

python 复制代码
ckpt/checkpoint0149.pth

b) 使用以下命令推断图片中的关系:

python 复制代码
python inference.py --img_path $IMAGE_PATH --resume $MODEL_PATH

我们附带了 VG 数据集的 5 张图片和 1 张互联网图片。您也可以用您自定义的图片进行测试。结果应该如下所示:

训练

a) 在单节点上使用 8 个 GPU 在 Visual Genome 上训练 RelTR(每个 GPU 2 张图片):

python 复制代码
python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --dataset vg --img_folder data/vg/images/ --ann_path data/vg/ --batch_size 2 --output_dir ckpt

b) 在单节点上使用 8 个 GPU 在 Open Images V6 上训练 RelTR(每个 GPU 2 张图片):

python 复制代码
python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --dataset oi --img_folder data/oi/images/ --ann_path data/oi/ --batch_size 2 --output_dir ckpt

评估

a) 使用单个 GPU 评估 Visual Genome 上的预训练 RelTR(每个 GPU 1 张图片):

python 复制代码
python main.py --dataset vg --img_folder data/vg/images/ --ann_path data/vg/ --eval --batch_size 1 --resume ckpt/checkpoint0149.pth

b) 使用单个 GPU 评估 Open Images V6 上的预训练 RelTR(每个 GPU 1 张图片):

python 复制代码
python main.py --dataset oi --img_folder data/oi/images/ --ann_path data/oi/ --eval --batch_size 1 --resume ckpt/checkpoint0149_oi.pth
相关推荐
杰尼橙子1 分钟前
深度解读Karpathy说的Software 3.0时代,感觉是个人的机会很大的时代呀
人工智能·openai
我爱一条柴ya25 分钟前
【AI大模型】线性回归:经典算法的深度解析与实战指南
人工智能·python·算法·ai·ai编程
Qiuner30 分钟前
【源力觉醒 创作者计划】开源、易用、强中文:文心一言4.5或是 普通人/非AI程序员 的第一款中文AI?
人工智能·百度·开源·文心一言·gitcode
未来之窗软件服务42 分钟前
chrome webdrive异常处理-session not created falled opening key——仙盟创梦IDE
前端·人工智能·chrome·仙盟创梦ide·东方仙盟·数据调式
AI街潜水的八角1 小时前
深度学习图像分类数据集—蘑菇识别分类
人工智能·深度学习·分类
飞睿科技1 小时前
乐鑫代理商飞睿科技,2025年AI智能语音助手市场发展趋势与乐鑫芯片解决方案分析
人工智能
许泽宇的技术分享1 小时前
从新闻到知识图谱:用大模型和知识工程“八步成诗”打造科技并购大脑
人工智能·科技·知识图谱
坤坤爱学习2.02 小时前
求医十年,病因不明,ChatGPT:你看起来有基因突变
人工智能·ai·chatgpt·程序员·大模型·ai编程·大模型学
蹦蹦跳跳真可爱5892 小时前
Python----循环神经网络(Transformer ----注意力机制)
人工智能·深度学习·nlp·transformer·循环神经网络
空中湖4 小时前
tensorflow武林志第二卷第九章:玄功九转
人工智能·python·tensorflow