File size: 3,167 Bytes
938e515 |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
// Copyright (c) Facebook, Inc. and its affiliates.
#include <torch/extension.h>
#include "ROIAlignRotated/ROIAlignRotated.h"
#include "box_iou_rotated/box_iou_rotated.h"
#include "cocoeval/cocoeval.h"
#include "deformable/deform_conv.h"
#include "nms_rotated/nms_rotated.h"
namespace detectron2 {
#if defined(WITH_CUDA) || defined(WITH_HIP)
extern int get_cudart_version();
#endif
std::string get_cuda_version() {
#if defined(WITH_CUDA) || defined(WITH_HIP)
std::ostringstream oss;
#if defined(WITH_CUDA)
oss << "CUDA ";
#else
oss << "HIP ";
#endif
// copied from
// https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/cuda/detail/CUDAHooks.cpp#L231
auto printCudaStyleVersion = [&](int v) {
oss << (v / 1000) << "." << (v / 10 % 100);
if (v % 10 != 0) {
oss << "." << (v % 10);
}
};
printCudaStyleVersion(get_cudart_version());
return oss.str();
#else // neither CUDA nor HIP
return std::string("not available");
#endif
}
bool has_cuda() {
#if defined(WITH_CUDA)
return true;
#else
return false;
#endif
}
// similar to
// https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Version.cpp
std::string get_compiler_version() {
std::ostringstream ss;
#if defined(__GNUC__)
#ifndef __clang__
#if ((__GNUC__ <= 4) && (__GNUC_MINOR__ <= 8))
#error "GCC >= 4.9 is required!"
#endif
{ ss << "GCC " << __GNUC__ << "." << __GNUC_MINOR__; }
#endif
#endif
#if defined(__clang_major__)
{
ss << "clang " << __clang_major__ << "." << __clang_minor__ << "."
<< __clang_patchlevel__;
}
#endif
#if defined(_MSC_VER)
{ ss << "MSVC " << _MSC_FULL_VER; }
#endif
return ss.str();
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("get_compiler_version", &get_compiler_version, "get_compiler_version");
m.def("get_cuda_version", &get_cuda_version, "get_cuda_version");
m.def("has_cuda", &has_cuda, "has_cuda");
m.def("deform_conv_forward", &deform_conv_forward, "deform_conv_forward");
m.def(
"deform_conv_backward_input",
&deform_conv_backward_input,
"deform_conv_backward_input");
m.def(
"deform_conv_backward_filter",
&deform_conv_backward_filter,
"deform_conv_backward_filter");
m.def(
"modulated_deform_conv_forward",
&modulated_deform_conv_forward,
"modulated_deform_conv_forward");
m.def(
"modulated_deform_conv_backward",
&modulated_deform_conv_backward,
"modulated_deform_conv_backward");
m.def("COCOevalAccumulate", &COCOeval::Accumulate, "COCOeval::Accumulate");
m.def(
"COCOevalEvaluateImages",
&COCOeval::EvaluateImages,
"COCOeval::EvaluateImages");
pybind11::class_<COCOeval::InstanceAnnotation>(m, "InstanceAnnotation")
.def(pybind11::init<uint64_t, double, double, bool, bool>());
pybind11::class_<COCOeval::ImageEvaluation>(m, "ImageEvaluation")
.def(pybind11::init<>());
}
TORCH_LIBRARY(detectron2, m) {
m.def("nms_rotated", &nms_rotated);
m.def("box_iou_rotated", &box_iou_rotated);
m.def("roi_align_rotated_forward", &ROIAlignRotated_forward);
m.def("roi_align_rotated_backward", &ROIAlignRotated_backward);
}
} // namespace detectron2
|