From f677ea31db8f45dbfec2fe5e519da82853815776 Mon Sep 17 00:00:00 2001 From: Francisco Massa Date: Wed, 18 Sep 2019 18:20:08 -0300 Subject: [PATCH] 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 | 22 +-- test/test_ops.py | 8 +- torchvision/__init__.py | 2 + torchvision/csrc/ROIAlign.h | 71 ++++++++ torchvision/csrc/ROIPool.h | 64 ++++++- torchvision/csrc/custom_ops/custom_ops.cpp | 159 ------------------ torchvision/csrc/vision.cpp | 46 +++-- torchvision/extension.py | 40 +++-- torchvision/ops/__init__.py | 4 + .../{_custom_ops.py => _register_onnx_ops.py} | 13 +- torchvision/ops/_utils.py | 18 +- torchvision/ops/boxes.py | 2 - torchvision/ops/roi_align.py | 10 +- torchvision/ops/roi_pool.py | 10 +- 14 files changed, 230 insertions(+), 239 deletions(-) delete mode 100644 torchvision/csrc/custom_ops/custom_ops.cpp rename torchvision/ops/{_custom_ops.py => _register_onnx_ops.py} (87%) diff --git a/setup.py b/setup.py index e9b4627ec77..74d39697513 100644 --- a/setup.py +++ b/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 diff --git a/test/test_ops.py b/test/test_ops.py index 7db8c6981d0..b7d41e8e6c3 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -190,7 +190,7 @@ def func(input): @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 @@ def func(input): @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 @@ def func(input): @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 @@ def func(input): @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' diff --git a/torchvision/__init__.py b/torchvision/__init__.py index 297aca2b228..84dbe4fa1ee 100644 --- a/torchvision/__init__.py +++ b/torchvision/__init__.py @@ -5,6 +5,8 @@ from torchvision import utils from torchvision import io +from .extension import _HAS_OPS + try: from .version import __version__ # noqa: F401 except ImportError: diff --git a/torchvision/csrc/ROIAlign.h b/torchvision/csrc/ROIAlign.h index 7e18cf68f57..765d4879d99 100644 --- a/torchvision/csrc/ROIAlign.h +++ b/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 { + 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]; +} diff --git a/torchvision/csrc/ROIPool.h b/torchvision/csrc/ROIPool.h index 7aefcc5e810..79b40293176 100644 --- a/torchvision/csrc/ROIPool.h +++ b/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 { + 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 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(result[0], result[1]); +} diff --git a/torchvision/csrc/custom_ops/custom_ops.cpp b/torchvision/csrc/custom_ops/custom_ops.cpp deleted file mode 100644 index e3b7bc9f0f0..00000000000 --- a/torchvision/csrc/custom_ops/custom_ops.cpp +++ /dev/null @@ -1,159 +0,0 @@ -#include -#include - -#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 { - 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 { - 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 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(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); diff --git a/torchvision/csrc/vision.cpp b/torchvision/csrc/vision.cpp index 61a4eeee727..243e2e87ff9 100644 --- a/torchvision/csrc/vision.cpp +++ b/torchvision/csrc/vision.cpp @@ -1,20 +1,44 @@ +#include +#include + +#ifdef WITH_CUDA +#include +#endif + #include "ROIAlign.h" #include "ROIPool.h" #include "nms.h" -#ifdef WITH_CUDA -#include +// 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); diff --git a/torchvision/extension.py b/torchvision/extension.py index b872af6ee74..790bca0d1f4 100644 --- a/torchvision/extension.py +++ b/torchvision/extension.py @@ -1,19 +1,34 @@ -_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() diff --git a/torchvision/ops/__init__.py b/torchvision/ops/__init__.py index fbd1181929b..a05c754989f 100644 --- a/torchvision/ops/__init__.py +++ b/torchvision/ops/__init__.py @@ -4,6 +4,10 @@ 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', diff --git a/torchvision/ops/_custom_ops.py b/torchvision/ops/_register_onnx_ops.py similarity index 87% rename from torchvision/ops/_custom_ops.py rename to torchvision/ops/_register_onnx_ops.py index 70108bdf152..cc30ad81db6 100644 --- a/torchvision/ops/_custom_ops.py +++ b/torchvision/ops/_register_onnx_ops.py @@ -1,16 +1,8 @@ -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 roi_pool(g, input, rois, spatial_scale, pooled_height, pooled_width): 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() diff --git a/torchvision/ops/_utils.py b/torchvision/ops/_utils.py index 67f0ea4feeb..269abaf7db3 100644 --- a/torchvision/ops/_utils.py +++ b/torchvision/ops/_utils.py @@ -1,24 +1,26 @@ 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 diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py index 3e773a02f8b..239a2446e22 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -1,5 +1,4 @@ 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) diff --git a/torchvision/ops/roi_align.py b/torchvision/ops/roi_align.py index 3038fb0dca0..abba99d420a 100644 --- a/torchvision/ops/roi_align.py +++ b/torchvision/ops/roi_align.py @@ -1,16 +1,14 @@ 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) diff --git a/torchvision/ops/roi_pool.py b/torchvision/ops/roi_pool.py index 6a9eaf6fdd9..50381c4ff2f 100644 --- a/torchvision/ops/roi_pool.py +++ b/torchvision/ops/roi_pool.py @@ -1,16 +1,14 @@ 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