|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
import torch |
|
|
|
from torch.utils.ffi import create_extension |
|
|
|
headers = [] |
|
sources = [] |
|
defines = [] |
|
extra_objects = [] |
|
with_cuda = False |
|
|
|
if torch.cuda.is_available(): |
|
with_cuda = True |
|
|
|
headers+= ['src/prroi_pooling_gpu.h'] |
|
sources += ['src/prroi_pooling_gpu.c'] |
|
defines += [('WITH_CUDA', None)] |
|
|
|
this_file = os.path.dirname(os.path.realpath(__file__)) |
|
extra_objects_cuda = ['src/prroi_pooling_gpu_impl.cu.o'] |
|
extra_objects_cuda = [os.path.join(this_file, fname) for fname in extra_objects_cuda] |
|
extra_objects.extend(extra_objects_cuda) |
|
else: |
|
|
|
raise NotImplementedError('Precise RoI Pooling only supports GPU (cuda) implememtations.') |
|
|
|
ffi = create_extension( |
|
'_prroi_pooling', |
|
headers=headers, |
|
sources=sources, |
|
define_macros=defines, |
|
relative_to=__file__, |
|
with_cuda=with_cuda, |
|
extra_objects=extra_objects |
|
) |
|
|
|
if __name__ == '__main__': |
|
ffi.build() |
|
|
|
|