Fix the python package generation with pytorch ops

isl-org · Jan 31, 2025 · 57ff4a2 · 57ff4a2
1 parent ef8df4b
commit 57ff4a2
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 4 deletions.
diff --git a/cpp/open3d/ml/pytorch/CMakeLists.txt b/cpp/open3d/ml/pytorch/CMakeLists.txt
@@ -123,7 +123,8 @@ set(TORCH_OPS_ARCH_DIR
     "${TORCH_OPS_DIR}/$<IF:$<BOOL:${BUILD_CUDA_MODULE}>,cuda,cpu>")
 set_target_properties(open3d_torch_ops PROPERTIES
     LIBRARY_OUTPUT_DIRECTORY "${TORCH_OPS_ARCH_DIR}"
-    ARCHIVE_OUTPUT_DIRECTORY "${TORCH_OPS_ARCH_DIR}")
+    ARCHIVE_OUTPUT_DIRECTORY "${TORCH_OPS_ARCH_DIR}"
+    RUNTIME_OUTPUT_DIRECTORY "${TORCH_OPS_ARCH_DIR}")
 
 # Do not add "lib" prefix
 set_target_properties(open3d_torch_ops PROPERTIES PREFIX "")

diff --git a/cpp/pybind/CMakeLists.txt b/cpp/pybind/CMakeLists.txt
@@ -169,8 +169,9 @@ endif()
 # add additional optional compiled modules
 if (BUILD_PYTORCH_OPS)
     list( APPEND COMPILED_MODULE_PATH_LIST $<TARGET_FILE:open3d_torch_ops> )
+    set(DEPENDENCIES_DIR "$<TARGET_FILE_DIR:tbb>")
     add_custom_command( OUTPUT "${CMAKE_BINARY_DIR}/lib/ml/torch/python/ops.py" "${CMAKE_BINARY_DIR}/lib/ml/torch/python/return_types.py"
-            COMMAND ${Python3_EXECUTABLE} generate_torch_ops_wrapper.py --input_ops_py_in "${PYTHON_PACKAGE_SRC_DIR}/open3d/ml/torch/python/ops.py.in" --input_return_types_py_in "${PYTHON_PACKAGE_SRC_DIR}/open3d/ml/torch/python/return_types.py.in" --output_dir "${CMAKE_BINARY_DIR}/lib/ml/torch/python/" --lib $<TARGET_FILE:open3d_torch_ops> --tensorflow_ops_dir "${CMAKE_CURRENT_SOURCE_DIR}/../open3d/ml/tensorflow"
+            COMMAND ${Python3_EXECUTABLE} generate_torch_ops_wrapper.py --input_ops_py_in "${PYTHON_PACKAGE_SRC_DIR}/open3d/ml/torch/python/ops.py.in" --input_return_types_py_in "${PYTHON_PACKAGE_SRC_DIR}/open3d/ml/torch/python/return_types.py.in" --output_dir "${CMAKE_BINARY_DIR}/lib/ml/torch/python/" --lib $<TARGET_FILE:open3d_torch_ops> --tensorflow_ops_dir "${CMAKE_CURRENT_SOURCE_DIR}/../open3d/ml/tensorflow" --dependencies_dir "${DEPENDENCIES_DIR}"
                         DEPENDS open3d_torch_ops
                         WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
                         COMMENT "Generating python ops.py and return_types.py" )

diff --git a/cpp/pybind/generate_torch_ops_wrapper.py b/cpp/pybind/generate_torch_ops_wrapper.py
@@ -130,12 +130,22 @@ def main():
                         type=str,
                         required=True,
                         help="This is cpp/open3d/ml/tensorflow")
+    parser.add_argument("--dependencies_dir",
+                        type=str,
+                        required=True,
+                        help="PyTorch dependencies directory")
 
     args = parser.parse_args()
     print(args)
 
+    if sys.platform == "win32":
+        open3d_deps = os.add_dll_directory(args.dependencies_dir)
+
     torch.ops.load_library(args.lib)
 
+    if sys.platform == "win32":
+        open3d_deps.close()
+
     generated_function_strs = ''
     generated_namedtuple_strs = ''
     for schema in torch._C._jit_get_all_schemas():

diff --git a/python/open3d/ml/tf/__init__.py b/python/open3d/ml/tf/__init__.py
@@ -11,7 +11,7 @@
 from open3d import _build_config
 
 if not _build_config["Tensorflow_VERSION"]:
-    raise Exception('Open3D was not built with TensorFlow support!')
+    raise ImportError('Open3D was not built with TensorFlow support!')
 
 _o3d_tf_version = _build_config["Tensorflow_VERSION"].split('.')
 if _tf_version.split('.')[:2] != _o3d_tf_version[:2]:

diff --git a/python/open3d/ml/torch/__init__.py b/python/open3d/ml/torch/__init__.py
@@ -12,7 +12,7 @@
 from open3d import _build_config
 
 if not _build_config["Pytorch_VERSION"]:
-    raise Exception('Open3D was not built with PyTorch support!')
+    raise ImportError('Open3D was not built with PyTorch support!')
 _o3d_torch_version = _verp(_build_config["Pytorch_VERSION"])
 # Check match with PyTorch version, any patch level is OK
 if _verp(_torch.__version__).release[:2] != _o3d_torch_version.release[:2]: