pytorch调用cpp

首先创建需要用到的文件，最后得到目录结构为

├── cpp_ex
│   ├── __init__.py
│   ├── cpp_py.py
│   └── csrc
│       └── add.cpp
└── main.py

其中main.py为运行文件。cpp_ex是需要调用的文件的包，main.py通过调用cpp_py.py来实现对cpp文件add.cpp的调用，起到中介的作用。但是实际上可以在main.py中实现。这里是为了让结构更加清楚，将其分开了。

非Tensor类型

从非Tensor类型入手的主要原因是

add.cpp

#include <torch/extension.h>

int add(int a, int b){
    return 0;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("add", &add, "add demo");
}

这里主要做了三件事

本例是利用pytorch调用cpp文件，所以导入*<torch/extension.h>*文件
定义add函数
利用pybind11将函数暴露出去，是py文件可以对其进行调用

pybind11是python与cpp文件互相调用的一个包，如果感兴趣可以查看https://pybind11.readthedocs.io/en/stable/basics.html

cpp_py.py

这里为了介绍方便将cpp_py.py分成三个部分

包的导入

1
2
3

import os
import torch.nn as nn
from torch.utils.cpp_extension import load

cpp函数的导入

def load_cpp_ext(exe_name):
    root_dir = os.path.join(os.path.split(__file__)[0])
    exe_csrc = os.path.join(root_dir, "csrc")
    exe_path = os.path.join(exe_csrc, "_ext", exe_name)
    os.makedirs(exe_path, exist_ok=True)
    exe_sources = [
        os.path.join(exe_csrc, "{}.cpp".format(exe_name))
    ]
    exe = load(
        name=exe_name,
        sources=exe_sources
    )
    return exe
    
_add = load_cpp_ext('add')

在这一步主要是调用load函数，从cpp文件中得到一个py可以执行的函数。所以，进行的操作大多是文件检索。值得一提的是，name是cpp文件生成的可执行文件名，也是py文件调用的函数名，注意保持一致。同时，load还可以传入其他与包编译相关的参数，这个在后面会提到。

利用pytorch使py文件可以调用

class Add(nn.Module):
    def __init__(self):
        super(Add, self).__init__()

    def forward(self, x, y):
        z = _add.add(x, y)
        return z

利用pytorch对add函数进行调用的方式和定义模型一样

main.py

from cpp_ex.cpp_demo import Add

a = 1
b = 2
add_func = Add()
c = add_func(a, b)

print(c)

Output:0

直接调用即可

pytorch与C++

pytorch调用cpp

非Tensor类型

add.cpp

cpp_py.py

main.py