pytorch使用c++/cuda扩展

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 addition
    return 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,CUDAExtension



setup(
    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.py

a = 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

参考

https://zhuanlan.zhihu.com/p/459955492

深入解析PyTorch中的PYBIND11_MODULE:功能与实现_pytorch pybind11-CSDN博客

相关推荐
CodeCraft Studio34 分钟前
CAD文件处理控件Aspose.CAD教程:使用 Python 将绘图转换为 Photoshop
python·photoshop·cad·aspose·aspose.cad
Python×CATIA工业智造3 小时前
Frida RPC高级应用:动态模拟执行Android so文件实战指南
开发语言·python·pycharm
onceco3 小时前
领域LLM九讲——第5讲 为什么选择OpenManus而不是QwenAgent(附LLM免费api邀请码)
人工智能·python·深度学习·语言模型·自然语言处理·自动化
狐凄4 小时前
Python实例题:基于 Python 的简单聊天机器人
开发语言·python
悦悦子a啊5 小时前
Python之--基本知识
开发语言·前端·python
笑稀了的野生俊6 小时前
在服务器中下载 HuggingFace 模型:终极指南
linux·服务器·python·bash·gpu算力
Naiva6 小时前
【小技巧】Python+PyCharm IDE 配置解释器出错,环境配置不完整或不兼容。(小智AI、MCP、聚合数据、实时新闻查询、NBA赛事查询)
ide·python·pycharm
路来了7 小时前
Python小工具之PDF合并
开发语言·windows·python
蓝婷儿7 小时前
Python 机器学习核心入门与实战进阶 Day 3 - 决策树 & 随机森林模型实战
人工智能·python·机器学习
AntBlack7 小时前
拖了五个月 ,不当韭菜体验版算是正式发布了
前端·后端·python