Revert "Register Torchvision Ops as Cutom Ops (#1267)" (#1316)

This reverts commit 78f169b.

Author: ezyang

.travis.yml

@@ -47,10 +47,6 @@ before_install:
   - pip install future
   - pip install pytest pytest-cov codecov
   - pip install mock
-  - |
-    if [[ $TRAVIS_PYTHON_VERSION == 3.6 ]]; then
-      pip install onnxruntime
-    fi
   - conda install av -c conda-forge
 
 

setup.py

@@ -96,21 +96,12 @@ def get_extensions():
     source_models = [os.path.join(models_dir, s) for s in source_models]
     tests = test_file + source_models
 
-    custom_ops_sources = [os.path.join(extensions_dir, "custom_ops", "custom_ops.cpp"),
-                          os.path.join(extensions_dir, "cpu", "nms_cpu.cpp"),
-                          os.path.join(extensions_dir, "cpu", "ROIAlign_cpu.cpp"),
-                          os.path.join(extensions_dir, "cpu", "ROIPool_cpu.cpp")]
-    custom_ops_sources_cuda = [os.path.join(extensions_dir, "cuda", "nms_cuda.cu"),
-                               os.path.join(extensions_dir, "cuda", "ROIAlign_cuda.cu"),
-                               os.path.join(extensions_dir, "cuda", "ROIPool_cuda.cu")]
-
     define_macros = []
 
     extra_compile_args = {}
     if (torch.cuda.is_available() and CUDA_HOME is not None) or os.getenv('FORCE_CUDA', '0') == '1':
         extension = CUDAExtension
         sources += source_cuda
-        custom_ops_sources += custom_ops_sources_cuda
         define_macros += [('WITH_CUDA', None)]
         nvcc_flags = os.getenv('NVCC_FLAGS', '')
         if nvcc_flags == '':
@@ -147,14 +138,7 @@ def get_extensions():
             include_dirs=tests_include_dirs,
             define_macros=define_macros,
             extra_compile_args=extra_compile_args,
-        ),
-        extension(
-            "torchvision._custom_ops",
-            sources=custom_ops_sources,
-            include_dirs=include_dirs,
-            define_macros=define_macros,
-            extra_compile_args=extra_compile_args,
-        ),
+        )
     ]
 
     return ext_modules
@@ -195,6 +179,5 @@ def run(self):
         "scipy": ["scipy"],
     },
     ext_modules=get_extensions(),
-    cmdclass={'build_ext': torch.utils.cpp_extension.BuildExtension,
-              'clean': clean}
+    cmdclass={'build_ext': torch.utils.cpp_extension.BuildExtension, 'clean': clean}
 )

test/test_onnx.py

@@ -10,11 +10,11 @@
 at::Tensor ROIAlign_forward(
     const at::Tensor& input, // Input feature map.
     const at::Tensor& rois, // List of ROIs to pool over.
-    const double spatial_scale, // The scale of the image features. ROIs will be
+    const float spatial_scale, // The scale of the image features. ROIs will be
     // scaled to this.
-    const int64_t pooled_height, // The height of the pooled feature map.
-    const int64_t pooled_width, // The width of the pooled feature
-    const int64_t sampling_ratio) // The number of points to sample in each bin
+    const int pooled_height, // The height of the pooled feature map.
+    const int pooled_width, // The width of the pooled feature
+    const int sampling_ratio) // The number of points to sample in each bin
 // along each axis.
 {
   if (input.type().is_cuda()) {

torchvision/csrc/ROIAlign.h

@@ -9,9 +9,9 @@
 std::tuple<at::Tensor, at::Tensor> ROIPool_forward(
     const at::Tensor& input,
     const at::Tensor& rois,
-    const double spatial_scale,
-    const int64_t pooled_height,
-    const int64_t pooled_width) {
+    const float spatial_scale,
+    const int pooled_height,
+    const int pooled_width) {
   if (input.type().is_cuda()) {
 #ifdef WITH_CUDA
     return ROIPool_forward_cuda(

torchvision/csrc/ROIPool.h

@@ -8,7 +8,7 @@
 at::Tensor nms(
     const at::Tensor& dets,
     const at::Tensor& scores,
-    const double iou_threshold) {
+    const float iou_threshold) {
   if (dets.device().is_cuda()) {
 #ifdef WITH_CUDA
     if (dets.numel() == 0) {

torchvision/csrc/custom_ops/custom_ops.cpp

@@ -7,8 +7,6 @@
 #endif
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  // TODO: remove nms from here since it is now registered
-  //       and used as a PyTorch custom op
   m.def("nms", &nms, "non-maximum suppression");
   m.def("roi_align_forward", &ROIAlign_forward, "ROIAlign_forward");
   m.def("roi_align_backward", &ROIAlign_backward, "ROIAlign_backward");

torchvision/csrc/nms.h

@@ -1,5 +1,5 @@
 import torch
-import torchvision.ops._custom_ops
+from torchvision.extension import _lazy_import
 
 
 def nms(boxes, scores, iou_threshold):
@@ -29,7 +29,8 @@ def nms(boxes, scores, iou_threshold):
         of the elements that have been kept
         by NMS, sorted in decreasing order of scores
     """
-    return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
+    _C = _lazy_import()
+    return _C.nms(boxes, scores, iou_threshold)
 
 
 def batched_nms(boxes, scores, idxs, iou_threshold):