当前位置：首页 > news >正文

pytorch使用c++/cuda扩展

news 来源：原创 2025/6/8 23:08:34

1、编写：c++/cuda拓展源文件

pybind11_demo/
├── setup.py
├── example.cpp
└── test.py

example.cpp

#include <torch/extension.h>
#include <vector>// Forward declaration of the function
torch::Tensor custom_add(torch::Tensor a, torch::Tensor b);// The actual implementation
torch::Tensor custom_add(torch::Tensor a, torch::Tensor b) {// Simple element-wise additionreturn a + b;
}// Pybind11 module definition
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("custom_add", &custom_add, "A function that adds two tensors");
}

PyTorch中的PYBIND11_MODULE

PYBIND11_MODULE是Pybind11库中的一个宏，它用于定义一个Python模块，并将C++类、函数或其他对象绑定到该模块。这使得Python可以直接调用C++编写的函数和类，极大地提高了Python的性能，尤其是当计算密集型任务需要底层C++实现时。

2、编译：setuptools指导c++/cuda拓展的编译

setup.py

from setuptools import setup, Extension
from torch.utils.cpp_extension import CppExtension, BuildExtension,CUDAExtensionsetup(name='python_demo', # python包的名称ext_modules=[CppExtension(name='demo', # 扩展模块名称，后面import使用sources=['example.cpp'],extra_compile_args={'CXX': ['-w', '-std=c++14']})],cmdclass={'build_ext': BuildExtension}
)# python setup.py install
# or for development:
# python setup.py develop

指定构建命令

cmdclass={  'build_ext': BuildExtension  
}

cmdclass是一个字典，用于指定自定义的构建命令。

'build_ext'是setuptools中的一个标准构建命令，用于构建扩展模块。

BuildExtension是PyTorch提供的BuildExtension类，它扩展了setuptools的build_ext命令，以支持C++和CUDA扩展的编译。

3、python调用编译完成的库

test.py

import torch
import demo  # The name you specified in setup.pya = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])result = demo.custom_add(a, b)
print(result)  # Should output tensor([5., 7., 9.])# python test.py