From 09ea8ba9fbe7a628e87dc10699bba64aaa7bc972 Mon Sep 17 00:00:00 2001 From: Soumith Chintala Date: Wed, 29 May 2019 21:56:22 -0700 Subject: [PATCH 1/2] make C extension lazy-import --- torchvision/__init__.py | 28 ---------------------------- torchvision/extension.py | 31 +++++++++++++++++++++++++++++++ torchvision/ops/boxes.py | 3 ++- torchvision/ops/roi_align.py | 4 +++- 4 files changed, 36 insertions(+), 30 deletions(-) create mode 100644 torchvision/extension.py diff --git a/torchvision/__init__.py b/torchvision/__init__.py index 4af0c08a8aa..82ba966dd5a 100644 --- a/torchvision/__init__.py +++ b/torchvision/__init__.py @@ -33,31 +33,3 @@ def get_image_backend(): Gets the name of the package used to load images """ return _image_backend - - -def _check_cuda_matches(): - """ - Make sure that CUDA versions match between the pytorch install and torchvision install - """ - import torch - from torchvision import _C - if hasattr(_C, "CUDA_VERSION") and torch.version.cuda is not None: - tv_version = str(_C.CUDA_VERSION) - if int(tv_version) < 10000: - tv_major = int(tv_version[0]) - tv_minor = int(tv_version[2]) - else: - tv_major = int(tv_version[0:2]) - tv_minor = int(tv_version[3]) - t_version = torch.version.cuda - t_version = t_version.split('.') - t_major = int(t_version[0]) - t_minor = int(t_version[1]) - if t_major != tv_major or t_minor != tv_minor: - raise RuntimeError("Detected that PyTorch and torchvision were compiled with different CUDA versions. " - "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)) - - -_check_cuda_matches() diff --git a/torchvision/extension.py b/torchvision/extension.py new file mode 100644 index 00000000000..06c41ceb5df --- /dev/null +++ b/torchvision/extension.py @@ -0,0 +1,31 @@ +_C = None + + +def _lazy_import(): + """ + Make sure that CUDA versions match between the pytorch install and torchvision install + """ + global _C + if _C is not None: + return _C + import torch + from torchvision import _C as C + _C = C + if hasattr(_C, "CUDA_VERSION") and torch.version.cuda is not None: + tv_version = str(_C.CUDA_VERSION) + if int(tv_version) < 10000: + tv_major = int(tv_version[0]) + tv_minor = int(tv_version[2]) + else: + tv_major = int(tv_version[0:2]) + tv_minor = int(tv_version[3]) + t_version = torch.version.cuda + t_version = t_version.split('.') + t_major = int(t_version[0]) + t_minor = int(t_version[1]) + if t_major != tv_major or t_minor != tv_minor: + raise RuntimeError("Detected that PyTorch and torchvision were compiled with different CUDA versions. " + "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 diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py index 897aa600fbf..5e9277ce80e 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -1,5 +1,5 @@ import torch -from torchvision import _C +from torchvision.extension import _lazy_import def nms(boxes, scores, iou_threshold): @@ -22,6 +22,7 @@ def nms(boxes, scores, iou_threshold): of the elements that have been kept by NMS, sorted in decreasing order of scores """ + _C = _lazy_import() return _C.nms(boxes, scores, iou_threshold) diff --git a/torchvision/ops/roi_align.py b/torchvision/ops/roi_align.py index 96d4f30323b..0302af4c478 100644 --- a/torchvision/ops/roi_align.py +++ b/torchvision/ops/roi_align.py @@ -6,7 +6,7 @@ from torch.nn.modules.utils import _pair -from torchvision import _C +from torchvision.extension import _lazy_import from ._utils import convert_boxes_to_roi_format @@ -18,6 +18,7 @@ def forward(ctx, input, roi, output_size, spatial_scale, sampling_ratio): ctx.spatial_scale = spatial_scale ctx.sampling_ratio = sampling_ratio ctx.input_shape = input.size() + _C = _lazy_import() output = _C.roi_align_forward( input, roi, spatial_scale, output_size[0], output_size[1], sampling_ratio) @@ -31,6 +32,7 @@ def backward(ctx, grad_output): spatial_scale = ctx.spatial_scale sampling_ratio = ctx.sampling_ratio bs, ch, h, w = ctx.input_shape + _C = _lazy_import() grad_input = _C.roi_align_backward( grad_output, rois, spatial_scale, output_size[0], output_size[1], bs, ch, h, w, sampling_ratio) From 6a75060b4b48ade0889d123bc9c22a6733ec9cb8 Mon Sep 17 00:00:00 2001 From: Soumith Chintala Date: Thu, 30 May 2019 07:01:33 -0700 Subject: [PATCH 2/2] add lazy loading to roi_pool --- torchvision/ops/roi_pool.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchvision/ops/roi_pool.py b/torchvision/ops/roi_pool.py index 3d8166517c1..dd78cf273fe 100644 --- a/torchvision/ops/roi_pool.py +++ b/torchvision/ops/roi_pool.py @@ -6,7 +6,7 @@ from torch.nn.modules.utils import _pair -from torchvision import _C +from torchvision.extension import _lazy_import from ._utils import convert_boxes_to_roi_format @@ -16,6 +16,7 @@ def forward(ctx, input, rois, output_size, spatial_scale): ctx.output_size = _pair(output_size) ctx.spatial_scale = spatial_scale ctx.input_shape = input.size() + _C = _lazy_import() output, argmax = _C.roi_pool_forward( input, rois, spatial_scale, output_size[0], output_size[1]) @@ -29,6 +30,7 @@ def backward(ctx, grad_output): output_size = ctx.output_size spatial_scale = ctx.spatial_scale bs, ch, h, w = ctx.input_shape + _C = _lazy_import() grad_input = _C.roi_pool_backward( grad_output, rois, argmax, spatial_scale, output_size[0], output_size[1], bs, ch, h, w)