diff --git a/torchvision/extension.py b/torchvision/extension.py index 69f34891837..ea837c234d3 100644 --- a/torchvision/extension.py +++ b/torchvision/extension.py @@ -1,3 +1,8 @@ +import ctypes +import os +import sys +from warnings import warn + import torch from ._internally_replaced_utils import _get_extension_path @@ -67,4 +72,22 @@ def _check_cuda_version(): return _version +def _load_library(lib_name): + lib_path = _get_extension_path(lib_name) + # On Windows Python-3.8+ has `os.add_dll_directory` call, + # which is called from _get_extension_path to configure dll search path + # Condition below adds a workaround for older versions by + # explicitly calling `LoadLibraryExW` with the following flags: + # - LOAD_LIBRARY_SEARCH_DEFAULT_DIRS (0x1000) + # - LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR (0x100) + if os.name == "nt" and sys.version_info < (3, 8): + _kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True) + if hasattr(_kernel32, "LoadLibraryExW"): + _kernel32.LoadLibraryExW(lib_path, None, 0x00001100) + else: + warn("LoadLibraryExW is missing in kernel32.dll") + + torch.ops.load_library(lib_path) + + _check_cuda_version() diff --git a/torchvision/io/_video_opt.py b/torchvision/io/_video_opt.py index 45bec44ec61..d75dd41534d 100644 --- a/torchvision/io/_video_opt.py +++ b/torchvision/io/_video_opt.py @@ -5,12 +5,11 @@ import torch -from .._internally_replaced_utils import _get_extension_path +from ..extension import _load_library try: - lib_path = _get_extension_path("video_reader") - torch.ops.load_library(lib_path) + _load_library("video_reader") _HAS_VIDEO_OPT = True except (ImportError, OSError): _HAS_VIDEO_OPT = False diff --git a/torchvision/io/image.py b/torchvision/io/image.py index f835565016c..68117c3b7ef 100644 --- a/torchvision/io/image.py +++ b/torchvision/io/image.py @@ -1,15 +1,15 @@ from enum import Enum +from warnings import warn import torch -from .._internally_replaced_utils import _get_extension_path +from ..extension import _load_library try: - lib_path = _get_extension_path("image") - torch.ops.load_library(lib_path) -except (ImportError, OSError): - pass + _load_library("image") +except (ImportError, OSError) as e: + warn(f"Failed to load image Python extension: {e}") class ImageReadMode(Enum):