Remove cpp extensions in favor of torch ops (#1348)

* Remove C++ extensions in favor of custom ops

* Remove unused custom_ops.cpp file

* Rename _custom_ops.py

* Reorganize functions

* Minor improvements and fixes

* Fix lint

* Fully scriptable ops

* Import types used by annotations

setup.py

@@ -52,9 +52,9 @@ def write_version_file():
     with open(version_path, 'w') as f:
         f.write("__version__ = '{}'\n".format(version))
         f.write("git_version = {}\n".format(repr(sha)))
-        f.write("from torchvision import _C\n")
-        f.write("if hasattr(_C, 'CUDA_VERSION'):\n")
-        f.write("    cuda = _C.CUDA_VERSION\n")
+        f.write("from torchvision.extension import _check_cuda_version\n")
+        f.write("if _check_cuda_version() > 0:\n")
+        f.write("    cuda = _check_cuda_version()\n")
 
 
 write_version_file()
@@ -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 == '':
@@ -148,13 +139,6 @@ def get_extensions():
             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

test/test_ops.py

@@ -190,7 +190,7 @@ class RoIPoolTester(unittest.TestCase):
 
         @torch.jit.script
         def script_func(input, rois):
-            return torch.ops.torchvision.roi_pool(input, rois, 1.0, 5, 5)[0]
+            return ops.roi_pool(input, rois, 5, 1.0)[0]
 
         assert gradcheck(lambda x: script_func(x, rois), (x,)), 'gradcheck failed for scripted roi_pool'
 
@@ -282,7 +282,7 @@ class RoIPoolTester(unittest.TestCase):
 
         @torch.jit.script
         def script_func(input, rois):
-            return torch.ops.torchvision.roi_pool(input, rois, 1.0, 5, 5)[0]
+            return ops.roi_pool(input, rois, 5, 1.0)[0]
 
         assert gradcheck(lambda x: script_func(x, rois), (x,)), 'gradcheck failed for scripted roi_pool on CUDA'
 
@@ -442,7 +442,7 @@ class RoIAlignTester(unittest.TestCase):
 
         @torch.jit.script
         def script_func(input, rois):
-            return torch.ops.torchvision.roi_align(input, rois, 0.5, 5, 5, 1)[0]
+            return ops.roi_align(input, rois, 5, 0.5, 1)[0]
 
         assert gradcheck(lambda x: script_func(x, rois), (x,)), 'gradcheck failed for scripted roi_align'
 
@@ -482,7 +482,7 @@ class RoIAlignTester(unittest.TestCase):
 
         @torch.jit.script
         def script_func(input, rois):
-            return torch.ops.torchvision.roi_align(input, rois, 0.5, 5, 5, 1)[0]
+            return ops.roi_align(input, rois, 5, 0.5, 1)[0]
 
         assert gradcheck(lambda x: script_func(x, rois), (x,)), 'gradcheck failed for scripted roi_align on CUDA'
 

torchvision/__init__.py

@@ -5,6 +5,8 @@ from torchvision import transforms
 from torchvision import utils
 from torchvision import io
 
+from .extension import _HAS_OPS
+
 try:
     from .version import __version__  # noqa: F401
 except ImportError:

torchvision/csrc/ROIAlign.h

@@ -74,3 +74,74 @@ at::Tensor ROIAlign_backward(
       width,
       sampling_ratio);
 }
+
+using namespace at;
+using torch::Tensor;
+using torch::autograd::AutogradContext;
+using torch::autograd::Variable;
+using torch::autograd::variable_list;
+
+class ROIAlignFunction : public torch::autograd::Function<ROIAlignFunction> {
+ public:
+  static variable_list forward(
+      AutogradContext* ctx,
+      Variable input,
+      Variable rois,
+      const double spatial_scale,
+      const int64_t pooled_height,
+      const int64_t pooled_width,
+      const int64_t sampling_ratio) {
+    ctx->saved_data["spatial_scale"] = spatial_scale;
+    ctx->saved_data["pooled_height"] = pooled_height;
+    ctx->saved_data["pooled_width"] = pooled_width;
+    ctx->saved_data["sampling_ratio"] = sampling_ratio;
+    ctx->saved_data["input_shape"] = input.sizes();
+    ctx->save_for_backward({rois});
+    auto result = ROIAlign_forward(
+        input,
+        rois,
+        spatial_scale,
+        pooled_height,
+        pooled_width,
+        sampling_ratio);
+    return {result};
+  }
+
+  static variable_list backward(
+      AutogradContext* ctx,
+      variable_list grad_output) {
+    // Use data saved in forward
+    auto saved = ctx->get_saved_variables();
+    auto rois = saved[0];
+    auto input_shape = ctx->saved_data["input_shape"].toIntList();
+    auto grad_in = ROIAlign_backward(
+        grad_output[0],
+        rois,
+        ctx->saved_data["spatial_scale"].toDouble(),
+        ctx->saved_data["pooled_height"].toInt(),
+        ctx->saved_data["pooled_width"].toInt(),
+        input_shape[0],
+        input_shape[1],
+        input_shape[2],
+        input_shape[3],
+        ctx->saved_data["sampling_ratio"].toInt());
+    return {
+        grad_in, Variable(), Variable(), Variable(), Variable(), Variable()};
+  }
+};
+
+Tensor roi_align(
+    const Tensor& input,
+    const Tensor& rois,
+    const double spatial_scale,
+    const int64_t pooled_height,
+    const int64_t pooled_width,
+    const int64_t sampling_ratio) {
+  return ROIAlignFunction::apply(
+      input,
+      rois,
+      spatial_scale,
+      pooled_height,
+      pooled_width,
+      sampling_ratio)[0];
+}

torchvision/csrc/ROIPool.h

@@ -63,4 +63,66 @@ at::Tensor ROIPool_backward(
       channels,
       height,
       width);
-}
+}
\ No newline at end of file
+
+using namespace at;
+using torch::Tensor;
+using torch::autograd::AutogradContext;
+using torch::autograd::Variable;
+using torch::autograd::variable_list;
+
+class ROIPoolFunction : public torch::autograd::Function<ROIPoolFunction> {
+ public:
+  static variable_list forward(
+      AutogradContext* ctx,
+      Variable input,
+      Variable rois,
+      const double spatial_scale,
+      const int64_t pooled_height,
+      const int64_t pooled_width) {
+    ctx->saved_data["spatial_scale"] = spatial_scale;
+    ctx->saved_data["pooled_height"] = pooled_height;
+    ctx->saved_data["pooled_width"] = pooled_width;
+    ctx->saved_data["input_shape"] = input.sizes();
+    auto result = ROIPool_forward(
+        input, rois, spatial_scale, pooled_height, pooled_width);
+    auto output = std::get<0>(result);
+    auto argmax = std::get<1>(result);
+    ctx->save_for_backward({rois, argmax});
+    ctx->mark_non_differentiable({argmax});
+    return {output, argmax};
+  }
+
+  static variable_list backward(
+      AutogradContext* ctx,
+      variable_list grad_output) {
+    // Use data saved in forward
+    auto saved = ctx->get_saved_variables();
+    auto rois = saved[0];
+    auto argmax = saved[1];
+    auto input_shape = ctx->saved_data["input_shape"].toIntList();
+    auto grad_in = ROIPool_backward(
+        grad_output[0],
+        rois,
+        argmax,
+        ctx->saved_data["spatial_scale"].toDouble(),
+        ctx->saved_data["pooled_height"].toInt(),
+        ctx->saved_data["pooled_width"].toInt(),
+        input_shape[0],
+        input_shape[1],
+        input_shape[2],
+        input_shape[3]);
+    return {grad_in, Variable(), Variable(), Variable(), Variable()};
+  }
+};
+
+std::tuple<Tensor, Tensor> roi_pool(
+    const Tensor& input,
+    const Tensor& rois,
+    const double spatial_scale,
+    const int64_t pooled_height,
+    const int64_t pooled_width) {
+  auto result = ROIPoolFunction::apply(
+      input, rois, spatial_scale, pooled_height, pooled_width);
+  return std::tuple<Tensor, Tensor>(result[0], result[1]);
+}

torchvision/csrc/custom_ops/custom_ops.cpp

-#include <Python.h>
-#include <torch/script.h>
-
-#include "ROIAlign.h"
-#include "ROIPool.h"
-#include "nms.h"
-
-using namespace at;
-
-// If we are in a Windows environment, we need to define
-// initialization functions for the _custom_ops extension
-#ifdef _WIN32
-#if PY_MAJOR_VERSION < 3
-PyMODINIT_FUNC init_custom_ops(void) {
-  // No need to do anything.
-  // _custom_ops.py will run on load
-  return NULL;
-}
-#else
-PyMODINIT_FUNC PyInit__custom_ops(void) {
-  // No need to do anything.
-  // _custom_ops.py will run on load
-  return NULL;
-}
-#endif
-#endif
-
-using torch::Tensor;
-using torch::autograd::AutogradContext;
-using torch::autograd::Variable;
-using torch::autograd::variable_list;
-
-class ROIAlignFunction : public torch::autograd::Function<ROIAlignFunction> {
- public:
-  static variable_list forward(
-      AutogradContext* ctx,
-      Variable input,
-      Variable rois,
-      const double spatial_scale,
-      const int64_t pooled_height,
-      const int64_t pooled_width,
-      const int64_t sampling_ratio) {
-    ctx->saved_data["spatial_scale"] = spatial_scale;
-    ctx->saved_data["pooled_height"] = pooled_height;
-    ctx->saved_data["pooled_width"] = pooled_width;
-    ctx->saved_data["sampling_ratio"] = sampling_ratio;
-    ctx->saved_data["input_shape"] = input.sizes();
-    ctx->save_for_backward({rois});
-    auto result = ROIAlign_forward(
-        input,
-        rois,
-        spatial_scale,
-        pooled_height,
-        pooled_width,
-        sampling_ratio);
-    return {result};
-  }
-
-  static variable_list backward(
-      AutogradContext* ctx,
-      variable_list grad_output) {
-    // Use data saved in forward
-    auto saved = ctx->get_saved_variables();
-    auto rois = saved[0];
-    auto input_shape = ctx->saved_data["input_shape"].toIntList();
-    auto grad_in = ROIAlign_backward(
-        grad_output[0],
-        rois,
-        ctx->saved_data["spatial_scale"].toDouble(),
-        ctx->saved_data["pooled_height"].toInt(),
-        ctx->saved_data["pooled_width"].toInt(),
-        input_shape[0],
-        input_shape[1],
-        input_shape[2],
-        input_shape[3],
-        ctx->saved_data["sampling_ratio"].toInt());
-    return {
-        grad_in, Variable(), Variable(), Variable(), Variable(), Variable()};
-  }
-};
-
-Tensor roi_align(
-    const Tensor& input,
-    const Tensor& rois,
-    const double spatial_scale,
-    const int64_t pooled_height,
-    const int64_t pooled_width,
-    const int64_t sampling_ratio) {
-  return ROIAlignFunction::apply(
-      input,
-      rois,
-      spatial_scale,
-      pooled_height,
-      pooled_width,
-      sampling_ratio)[0];
-}
-
-class ROIPoolFunction : public torch::autograd::Function<ROIPoolFunction> {
- public:
-  static variable_list forward(
-      AutogradContext* ctx,
-      Variable input,
-      Variable rois,
-      const double spatial_scale,
-      const int64_t pooled_height,
-      const int64_t pooled_width) {
-    ctx->saved_data["spatial_scale"] = spatial_scale;
-    ctx->saved_data["pooled_height"] = pooled_height;
-    ctx->saved_data["pooled_width"] = pooled_width;
-    ctx->saved_data["input_shape"] = input.sizes();
-    auto result = ROIPool_forward(
-        input, rois, spatial_scale, pooled_height, pooled_width);
-    auto output = std::get<0>(result);
-    auto argmax = std::get<1>(result);
-    ctx->save_for_backward({rois, argmax});
-    ctx->mark_non_differentiable({argmax});
-    return {output, argmax};
-  }
-
-  static variable_list backward(
-      AutogradContext* ctx,
-      variable_list grad_output) {
-    // Use data saved in forward
-    auto saved = ctx->get_saved_variables();
-    auto rois = saved[0];
-    auto argmax = saved[1];
-    auto input_shape = ctx->saved_data["input_shape"].toIntList();
-    auto grad_in = ROIPool_backward(
-        grad_output[0],
-        rois,
-        argmax,
-        ctx->saved_data["spatial_scale"].toDouble(),
-        ctx->saved_data["pooled_height"].toInt(),
-        ctx->saved_data["pooled_width"].toInt(),
-        input_shape[0],
-        input_shape[1],
-        input_shape[2],
-        input_shape[3]);
-    return {grad_in, Variable(), Variable(), Variable(), Variable()};
-  }
-};
-
-std::tuple<Tensor, Tensor> roi_pool(
-    const Tensor& input,
-    const Tensor& rois,
-    const double spatial_scale,
-    const int64_t pooled_height,
-    const int64_t pooled_width) {
-  auto result = ROIPoolFunction::apply(
-      input, rois, spatial_scale, pooled_height, pooled_width);
-  return std::tuple<Tensor, Tensor>(result[0], result[1]);
-}
-
-static auto registry =
-    torch::RegisterOperators()
-        .op("torchvision::nms", &nms)
-        .op("torchvision::roi_align(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width, int sampling_ratio) -> Tensor",
-            &roi_align)
-        .op("torchvision::roi_pool", &roi_pool);

torchvision/csrc/vision.cpp

+#include <Python.h>
+#include <torch/script.h>
+
+#ifdef WITH_CUDA
+#include <cuda.h>
+#endif
+
 #include "ROIAlign.h"
 #include "ROIPool.h"
 #include "nms.h"
 
-#ifdef WITH_CUDA
-#include <cuda.h>
+// If we are in a Windows environment, we need to define
+// initialization functions for the _custom_ops extension
+#ifdef _WIN32
+#if PY_MAJOR_VERSION < 3
+PyMODINIT_FUNC init_custom_ops(void) {
+  // No need to do anything.
+  // _custom_ops.py will run on load
+  return NULL;
+}
+#else
+PyMODINIT_FUNC PyInit__custom_ops(void) {
+  // No need to do anything.
+  // _custom_ops.py will run on load
+  return NULL;
+}
+#endif
 #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");
-  m.def("roi_pool_forward", &ROIPool_forward, "ROIPool_forward");
-  m.def("roi_pool_backward", &ROIPool_backward, "ROIPool_backward");
+int64_t _cuda_version() {
 #ifdef WITH_CUDA
-  m.attr("CUDA_VERSION") = CUDA_VERSION;
+  return CUDA_VERSION;
+#else
+  return -1;
 #endif
 }
+
+static auto registry =
+    torch::RegisterOperators()
+        .op("torchvision::nms", &nms)
+        .op("torchvision::roi_align(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width, int sampling_ratio) -> Tensor",
+            &roi_align)
+        .op("torchvision::roi_pool", &roi_pool)
+        .op("torchvision::_cuda_version", &_cuda_version);

torchvision/extension.py

-_C = None
+_HAS_OPS = False
 
 
-def _lazy_import():
+def _register_extensions():
+    import os
+    import imp
+    import torch
+
+    # load the custom_op_library and register the custom ops
+    lib_dir = os.path.dirname(__file__)
+    _, path, _ = imp.find_module("_C", [lib_dir])
+    torch.ops.load_library(path)
+
+
+try:
+    _register_extensions()
+    _HAS_OPS = True
+except (ImportError, OSError):
+    pass
+
+
+def _check_cuda_version():
     """
     Make sure that CUDA versions match between the pytorch install and torchvision install
     """
-    global _C
-    if _C is not None:
-        return _C
+    if not _HAS_OPS:
+        return -1
     import torch
-    from torchvision import _C as C
-    import torchvision.ops._custom_ops
-    _C = C
-    if hasattr(_C, "CUDA_VERSION") and torch.version.cuda is not None:
-        tv_version = str(_C.CUDA_VERSION)
+    _version = torch.ops.torchvision._cuda_version()
+    if _version != -1 and torch.version.cuda is not None:
+        tv_version = str(_version)
         if int(tv_version) < 10000:
             tv_major = int(tv_version[0])
             tv_minor = int(tv_version[2])
@@ -29,4 +44,7 @@ def _lazy_import():
                                "PyTorch has CUDA Version={}.{} and torchvision has CUDA Version={}.{}. "
                                "Please reinstall the torchvision that matches your PyTorch install."
                                .format(t_major, t_minor, tv_major, tv_minor))
-    return _C
+    return _version
+
+
+_check_cuda_version()

torchvision/ops/__init__.py

@@ -4,6 +4,10 @@ from .roi_pool import roi_pool, RoIPool
 from .poolers import MultiScaleRoIAlign
 from .feature_pyramid_network import FeaturePyramidNetwork
 
+from ._register_onnx_ops import _register_custom_op
+
+_register_custom_op()
+
 
 __all__ = [
     'nms', 'roi_align', 'RoIAlign', 'roi_pool', 'RoIPool',

torchvision/ops/_custom_ops.py → torchvision/ops/_register_onnx_ops.py

-import os
 import sys
-import imp
 import torch
 
 
-# load the custom_op_library and register the custom ops
-lib_dir = os.path.join(os.path.dirname(__file__), '..')
-file, path, description = imp.find_module("_custom_ops", [lib_dir])
-torch.ops.load_library(path)
-
-
-def register_custom_op():
+def _register_custom_op():
     from torch.onnx.symbolic_helper import parse_args, scalar_type_to_onnx
     from torch.onnx.symbolic_opset9 import select, unsqueeze, squeeze, _cast_Long, reshape
 
@@ -41,6 +33,3 @@ def register_custom_op():
     register_custom_op_symbolic('torchvision::nms', symbolic_multi_label_nms, 10)
     register_custom_op_symbolic('torchvision::roi_align', roi_align, 10)
     register_custom_op_symbolic('torchvision::roi_pool', roi_pool, 10)
-
-
-register_custom_op()

torchvision/ops/_utils.py

 import torch
+from torch import Tensor
+from torch.jit.annotations import List
 
 
 def _cat(tensors, dim=0):
+    # type: (List[Tensor], int) -> Tensor
     """
     Efficient version of torch.cat that avoids a copy if there is only a single element in a list
     """
-    assert isinstance(tensors, (list, tuple))
+    # TODO add back the assert
+    # assert isinstance(tensors, (list, tuple))
     if len(tensors) == 1:
         return tensors[0]
     return torch.cat(tensors, dim)
 
 
 def convert_boxes_to_roi_format(boxes):
+    # type: (List[Tensor]) -> Tensor
     concat_boxes = _cat([b for b in boxes], dim=0)
-    ids = _cat(
-        [
-            torch.full_like(b[:, :1], i)
-            for i, b in enumerate(boxes)
-        ],
-        dim=0,
-    )
+    temp = []
+    for i, b in enumerate(boxes):
+        temp.append(torch.full_like(b[:, :1], i))
+    ids = _cat(temp, dim=0)
     rois = torch.cat([ids, concat_boxes], dim=1)
     return rois

torchvision/ops/boxes.py

 import torch
-from torchvision.extension import _lazy_import
 
 
 def nms(boxes, scores, iou_threshold):
@@ -29,7 +28,6 @@ def nms(boxes, scores, iou_threshold):
         of the elements that have been kept
         by NMS, sorted in decreasing order of scores
     """
-    _lazy_import()
     return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
 
 

torchvision/ops/roi_align.py

 import torch
-from torch import nn
-
-from torch.autograd import Function
-from torch.autograd.function import once_differentiable
+from torch import nn, Tensor
 
 from torch.nn.modules.utils import _pair
+from torch.jit.annotations import List
 
-from torchvision.extension import _lazy_import
 from ._utils import convert_boxes_to_roi_format
 
 
 def roi_align(input, boxes, output_size, spatial_scale=1.0, sampling_ratio=-1):
+    # type: (Tensor, Tensor, int, float, int) -> Tensor
     """
     Performs Region of Interest (RoI) Align operator described in Mask R-CNN
 
@@ -35,9 +33,9 @@ def roi_align(input, boxes, output_size, spatial_scale=1.0, sampling_ratio=-1):
         output (Tensor[K, C, output_size[0], output_size[1]])
     """
     rois = boxes
+    output_size = _pair(output_size)
     if not isinstance(rois, torch.Tensor):
         rois = convert_boxes_to_roi_format(rois)
-    _lazy_import()
     return torch.ops.torchvision.roi_align(input, rois, spatial_scale,
                                            output_size[0], output_size[1],
                                            sampling_ratio)

torchvision/ops/roi_pool.py

 import torch
-from torch import nn
-
-from torch.autograd import Function
-from torch.autograd.function import once_differentiable
+from torch import nn, Tensor
 
 from torch.nn.modules.utils import _pair
+from torch.jit.annotations import List
 
-from torchvision.extension import _lazy_import
 from ._utils import convert_boxes_to_roi_format
 
 
 def roi_pool(input, boxes, output_size, spatial_scale=1.0):
+    # type: (Tensor, Tensor, int, float) -> Tensor
     """
     Performs Region of Interest (RoI) Pool operator described in Fast R-CNN
 
@@ -30,9 +28,9 @@ def roi_pool(input, boxes, output_size, spatial_scale=1.0):
         output (Tensor[K, C, output_size[0], output_size[1]])
     """
     rois = boxes
+    output_size = _pair(output_size)
     if not isinstance(rois, torch.Tensor):
         rois = convert_boxes_to_roi_format(rois)
-    _lazy_import()
     output, _ = torch.ops.torchvision.roi_pool(input, rois, spatial_scale,
                                                output_size[0], output_size[1])
     return output