Skip to content

GitLab

Unverified Commit f677ea31 authored by Francisco Massa's avatar Francisco Massa Committed by GitHub
Browse files

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
parent 0dd55882
Show whitespace changes
Inline Side-by-side
Showing with 230 additions and 239 deletions
setup.py
View file @ f677ea31
......@@ -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
View file @ f677ea31
......@@ -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
View file @ f677ea31
......@@ -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
View file @ f677ea31
......@@ -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
View file @ f677ea31
......@@ -64,3 +64,65 @@ at::Tensor ROIPool_backward(
height,
width);
}
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 deleted 100644 → 0
View file @ 0dd55882
#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
View file @ f677ea31
#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
View file @ f677ea31
_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
View file @ f677ea31
......@@ -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
View file @ f677ea31
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
View file @ f677ea31
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
View file @ f677ea31
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
View file @ f677ea31
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
View file @ f677ea31
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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment