Python调用自定义PyTorch C++扩展库失败问题排查求助

核心问题原因 #

你在setup.py中定义的C++扩展模块名为trans，但尝试从包self_sample_full中调用函数——这两个并非同一模块，因此找不到对应属性。此外还有CUDA张量访问的潜在问题，一并说明。

解决方案一：调整导入方式 #

直接导入编译后的扩展模块trans，修改use.py：

import torch
import trans

n = 10  # 补充缺失的n定义
output = torch.empty(n).to(device="cuda:0")
# 注意：当前函数仅能处理CPU张量，需先移回CPU
output_cpu = output.cpu()
trans.torch_tensor2vector(output_cpu)

解决方案二：将扩展模块归属到目标包 #

修改setup.py，让扩展模块成为self_sample_full的子模块：

from setuptools import setup, find_packages
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(
    name="self_sample_full",
    include_dirs=["include"],
    ext_modules=[
        # 修改扩展模块名为包的子模块
        CppExtension("self_sample_full.trans", ["pytorch/trans.cpp"]),
    ],
    cmdclass={
        "build_ext": BuildExtension
    },
    packages=find_packages()
)

同时需要在self_sample_full目录下创建空的__init__.py文件，使其成为合法Python包。

修改use.py：

import torch
from self_sample_full import trans

n = 10
output = torch.empty(n).to(device="cuda:0")
output_cpu = output.cpu()
trans.torch_tensor2vector(output_cpu)

额外注意事项 #

1. 安装命令：确保使用正确的安装方式，比如开发模式安装：pip install -e .，避免安装后模块未更新。
2. 路径检查：确认pytorch/trans.cpp路径相对于setup.py的位置正确，文件确实存在。
3. CUDA张量访问问题：你的torch_tensor2vector函数直接通过data_ptr()获取CPU内存指针，但传入的是CUDA张量，这会导致未定义行为。必须先将张量移到CPU（如上述代码中的output.cpu()），或修改C++代码使用CUDA API处理GPU内存。

内容的提问来源于stack exchange，提问作者YA xiang