[Windows] Workaround for loading bundled DLLs (#4893) (#5094)

Author: malfet

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
@@ -62,4 +67,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()

torchvision/io/_video_opt.py

@@ -8,12 +8,11 @@
 import numpy as np
 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

torchvision/io/image.py

@@ -1,14 +1,14 @@
 import torch
 from enum import Enum
+from warnings import warn
 
-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):