绪论

Pytorch中为什么会调用C++呢?目前来说,本人遇到的情况有两种:
为了

  1. 提高运算速度,C++偏向于底层,运算速度更快(相对于Python至少提升10倍);
  2. Pytorch中的内置函数不能满足自己想要实现的算法,不得不自己写code实现(我就是这种,好惨啊)。而且更重要的是,Pytorch的底层是C++(逐渐从c向C++转),所以你不得不学一手C++。

运行环境

个人的一点小体会,如果你刚开始上手学习深度学习,不管是Pytorch还是Tensorflow,在windows下学习运行就可以。但是,一旦你上复杂的项目或者实战,那必须到 Linux 系统。为什么呢?

  1. Linux对深度学习环境更加友好,某些命令在windows下根本无法编译;
  2. 作为一个学习人员,Github上的code大多是在Linux下编译和运行的,所以没办法了hxd,Linux欢迎你。

下面开始Pytorch中调用C++的正片!!!

先说明一下三个文件各自的用途:
setup.py : 用于关联 C++ 文件;
xxx.cpp : 用于存放自己实现的功能;
test.py :用于测试上述两个文件是否能正常使用。
这些都在同一个文件夹。
注意,在命令行输入这些时,不要忘了激活 Pytorch 的虚拟环境。

setup.py

from setuptools import setup, Extension
from torch.utils.cpp_extension import CppExtension, CUDAExtension, BuildExtension

"""
hello : 我们定义的名字,相当于Python的包名
hello.cpp : 就是下面的xxx.cpp文件,具体化
"""

setup(name='hello',
	  # 如果用到了cuda, 将下面的 CppExtension 换成 CUDAExtension
      ext_modules=[CppExtension('hello_cpp', ['hello.cpp'])],
      cmdclass={'build_ext': BuildExtension})

xxx.cpp

// 本文件名为hello.cpp
#include <iostream>
#include <torch/extension.h>

using namespace std;

int add(int i, int j){
    return i + j;
}

int mul(int i, int j){
    return i * j;
}


PYBIND11_MODULE(hello, m){
    m.def("add", &add, "int add");		// 第一个参数为python调用时的函数名,第二个参数为c++函数名,第三个为相应的描述
    m.def("mul", &mul, "int multiply");
}

在命令行将上述两个文件编译成 .so 文件:

python setup.py build_ext --inplace

test.py

import hello

a = hello.add(1, 2)
b = hello.mul(1, 2)
print("a = {}, b = {}".format(a,b))

结果:

a = 3, b = 2

但是在编译器中,例如Pycharm是这样的:

注意到,导入 hello 的时候,下面有红色下划线报错,但是它能够运行,看起来巨坑,谁能想到红了还能正常运行?

查了一些资料,我找到了解决的办法,就是把这个 hello 当包一样的安装在当前环境下,就像Numpy一样。

在当前路径的命令行输入:

python setup.py install

这样就不会报红了,看起来正常了。至此,完成了自己写Pytorch的包,当前你也可以这样扩展Python的包。

整理不易(弯路真的多),点个赞再走吧,期待与您共同进步。

更多推荐

Pytorch中调用C++(踩坑无数总结篇)