RA-SGG: Retrieval-Augmented Scene Graph Generation Framework via Multi-Prototype Learning

We refer to the implementation of PE-Net PENet.

Check INSTALL.md for installation instructions.

Dataset

Check DATASET.md for instructions of dataset preprocessing.

Train

ReTAG requires the pre-trained PE-Net and memory bank, which is populated with the relation embedding of training dataset.

Please download pre-trained models and the features for the memory bank

You can train ReTAG using scripts

bash scripts/predcls_train_retag.sh

Test

You can check the result of ReTAG in Model_Zoo.md

We provide scripts for testing the models

bash script/test.sh

Device

All our experiments are conducted on two NVIDIA GeForce RTX 3090 or one NVIDIA A6000, if you wanna run it on your own device, make sure to follow distributed training instructions in Scene-Graph-Benchmark.pytorch.

Tips

We use the rel_nms operation provided by RU-Net and HL-Net in PredCls and SGCls to filter the predicted relation predicates, which encourages diverse prediction results.

Acknowledgement

The code is implemented based on Scene-Graph-Benchmark.pytorch.