Registering operators in their files.

datumbox · datumbox · commit df0e8a7a7074 · 2020-11-30T21:32:52.000Z
diff --git a/torchvision/csrc/cpu/deform_conv2d_kernel.cpp b/torchvision/csrc/cpu/deform_conv2d_kernel.cpp
@@ -66,6 +66,8 @@
 // modified from
 // https://github.com/open-mmlab/mmdetection/blob/master/mmdet/ops/dcn/src/deform_conv_cuda.cpp
 
+#include <torch/script.h>
+
 #include "deform_conv2d_kernel.h"
 
 namespace {
@@ -1137,3 +1139,8 @@ deform_conv2d_backward_cpu(
   return std::make_tuple(
       grad_input, grad_weight, grad_offset, grad_mask, grad_bias);
 }
+
+TORCH_LIBRARY_IMPL(torchvision, CPU, m) {
+  m.impl("deform_conv2d", deform_conv2d_forward_cpu);
+  m.impl("_deform_conv2d_backward", deform_conv2d_backward_cpu);
+}
diff --git a/torchvision/csrc/cuda/deform_conv2d_kernel.cu b/torchvision/csrc/cuda/deform_conv2d_kernel.cu
@@ -70,6 +70,7 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <THC/THCAtomics.cuh>
+#include <torch/script.h>
 
 #include "cuda_helpers.h"
 #include "deform_conv2d_kernel.h"
@@ -1188,3 +1189,8 @@ deform_conv2d_backward_cuda(
   return std::make_tuple(
       grad_input, grad_weight, grad_offset, grad_mask, grad_bias);
 }
+
+TORCH_LIBRARY_IMPL(torchvision, CUDA, m) {
+  m.impl("deform_conv2d", deform_conv2d_forward_cuda);
+  m.impl("_deform_conv2d_backward", deform_conv2d_backward_cuda);
+}
diff --git a/torchvision/csrc/deform_conv2d.cpp b/torchvision/csrc/deform_conv2d.cpp
@@ -74,6 +74,10 @@ at::Tensor deform_conv2d_autocast(
              use_mask)
       .to(input.scalar_type());
 }
+
+TORCH_LIBRARY_IMPL(torchvision, Autocast, m) {
+  m.impl("deform_conv2d", deform_conv2d_autocast);
+}
 #endif
 
 std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
@@ -361,3 +365,8 @@ deform_conv2d_backward_autograd(
 
   return std::make_tuple(result[0], result[1], result[2], result[3], result[4]);
 }
+
+TORCH_LIBRARY_IMPL(torchvision, Autograd, m) {
+  m.impl("deform_conv2d", deform_conv2d_autograd);
+  m.impl("_deform_conv2d_backward", deform_conv2d_backward_autograd);
+}
diff --git a/torchvision/csrc/vision.cpp b/torchvision/csrc/vision.cpp
@@ -62,8 +62,6 @@ TORCH_LIBRARY(torchvision, m) {
 }
 
 TORCH_LIBRARY_IMPL(torchvision, CPU, m) {
-  m.impl("deform_conv2d", deform_conv2d_forward_cpu);
-  m.impl("_deform_conv2d_backward", deform_conv2d_backward_cpu);
   m.impl("nms", nms_cpu);
   m.impl("ps_roi_align", PSROIAlign_forward_cpu);
   m.impl("_ps_roi_align_backward", PSROIAlign_backward_cpu);
@@ -78,8 +76,6 @@ TORCH_LIBRARY_IMPL(torchvision, CPU, m) {
 // TODO: Place this in a hypothetical separate torchvision_cuda library
 #if defined(WITH_CUDA) || defined(WITH_HIP)
 TORCH_LIBRARY_IMPL(torchvision, CUDA, m) {
-  m.impl("deform_conv2d", deform_conv2d_forward_cuda);
-  m.impl("_deform_conv2d_backward", deform_conv2d_backward_cuda);
   m.impl("nms", nms_cuda);
   m.impl("ps_roi_align", PSROIAlign_forward_cuda);
   m.impl("_ps_roi_align_backward", PSROIAlign_backward_cuda);
@@ -95,7 +91,6 @@ TORCH_LIBRARY_IMPL(torchvision, CUDA, m) {
 // Autocast only needs to wrap forward pass ops.
 #if defined(WITH_CUDA) || defined(WITH_HIP)
 TORCH_LIBRARY_IMPL(torchvision, Autocast, m) {
-  m.impl("deform_conv2d", deform_conv2d_autocast);
   m.impl("nms", nms_autocast);
   m.impl("ps_roi_align", PSROIAlign_autocast);
   m.impl("ps_roi_pool", PSROIPool_autocast);
@@ -105,8 +100,6 @@ TORCH_LIBRARY_IMPL(torchvision, Autocast, m) {
 #endif
 
 TORCH_LIBRARY_IMPL(torchvision, Autograd, m) {
-  m.impl("deform_conv2d", deform_conv2d_autograd);
-  m.impl("_deform_conv2d_backward", deform_conv2d_backward_autograd);
   m.impl("ps_roi_align", PSROIAlign_autograd);
   m.impl("_ps_roi_align_backward", PSROIAlign_backward_autograd);
   m.impl("ps_roi_pool", PSROIPool_autograd);
