fast-reid/fastreid/evaluation/GPU-Re-Ranking/extension/propagation/setup.py

37 lines
1.2 KiB
Python

"""
Understanding Image Retrieval Re-Ranking: A Graph Neural Network Perspective
Xuanmeng Zhang, Minyue Jiang, Zhedong Zheng, Xiao Tan, Errui Ding, Yi Yang
Project Page : https://github.com/Xuanmeng-Zhang/gnn-re-ranking
Paper: https://arxiv.org/abs/2012.07620v2
======================================================================
On the Market-1501 dataset, we accelerate the re-ranking processing from 89.2s to 9.4ms
with one K40m GPU, facilitating the real-time post-processing. Similarly, we observe
that our method achieves comparable or even better retrieval results on the other four
image retrieval benchmarks, i.e., VeRi-776, Oxford-5k, Paris-6k and University-1652,
with limited time cost.
"""
from setuptools import setup, Extension
import torch
import torch.nn as nn
from torch.autograd import Function
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='gnn_propagate',
ext_modules=[
CUDAExtension('gnn_propagate', [
'gnn_propagate.cpp',
'gnn_propagate_kernel.cu',
]),
],
cmdclass={
'build_ext':BuildExtension
})