File size: 1,343 Bytes
92f0e98 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File : build.py
# Author : Jiayuan Mao, Tete Xiao
# Email : maojiayuan@gmail.com, jasonhsiao97@gmail.com
# Date : 07/13/2018
#
# This file is part of PreciseRoIPooling.
# Distributed under terms of the MIT license.
# Copyright (c) 2017 Megvii Technology Limited.
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:
# TODO(Jiayuan Mao @ 07/13): remove this restriction after we support the cpu implementation.
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()
|