Pytorch C Extension(C拓展)(即时编译JIT)(pybind11)(torch.utils.cpp_extension)

最新推荐文章于 2024-07-24 14:20:56 发布

hxxjxw

最新推荐文章于 2024-07-24 14:20:56 发布

阅读量1.6k

点赞数 1

本文链接：https://blog.csdn.net/hxxjxw/article/details/125462914

版权

本文介绍了如何在PyTorch环境中使用pybind11创建并即时编译C++扩展。通过torch.utils.cpp_extension.load()函数，可以便捷地将C++代码集成到PyTorch项目中，实现自定义运算。示例中展示了简单的加法操作，并展示了如何在Python中调用这个C++扩展函数。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在Pytorch中也可以自定义C拓展，并调用

并且可以不用setuptools, 采用即时编译(JIT)的方法来构建C++扩展。即时编译机制提供了一种简单调用PyTorch API中名为torch.utils.cpp_extension.load()的函数，来编译和加载扩展的方法：

这里还是采用pybind11的例子

hxx_test.cpp
#include <pybind11/pybind11.h>
 
int add( int i, int j ){
    return i+j;
}
 
PYBIND11_MODULE(python_example, m ){
    m.doc() = "pybind11 example";
    m.def("add", &add, "add two number" );
}
test.py
import torch
from torch.utils.cpp_extension import load

cpp_ext = load(
     name = "python_example",
     sources = ["hxx_test.cpp",], 
     verbose=True,
)

print(cpp_ext.add(2,3))
注意torch.utils.cpp_extension.load中的name参数一定要和PYBIND11_MODULE中的参数对应起来