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

pytorch拓展——Ubuntu vscode配置pytorch C++拓展环境

article 2025/1/16 22:00:14

引言：pytorch的官方教程中

例子1是用numpy+Funciton手动写的一个前向传播和反向传播
例子2是一个有需要优化参数的例子
例子3是用c++写的拓展，1.0之前是用c写的拓展
这里记录用vscode配置c++，然后导入torch的TH.h头文件的过程

ubuntu配置vscode支持c++调试

具体参考网上的教程: Ubuntu16.04下配置VScode的C/C++开发环境

注意：

tasks.json中需要加上“-std=c++11”这个参数表示支持c++11标准

重点：将pytorch的头文件TH.h等加入到C++的编译路径

1. 软链接的方式

找到c++编译器的路径，本人的在/usr/include/c++/9/, 可以通过一下命令找到位置：

locate iostream

找到pytorch的拓展TH.h所在的位置，命令：locate TH.h 或者find -name TH.h，如下是本人的：

/home/yunjie/.local/lib/python3.8/site-packages/torch/include/ATen/core/UnsafeFromTH.h
/home/yunjie/.local/lib/python3.8/site-packages/torch/include/TH/TH.h
/home/yunjie/anaconda3/lib/python3.10/site-packages/torch/include/ATen/core/UnsafeFromTH.h
/home/yunjie/anaconda3/pkgs/pytorch-1.12.1-cpu_py310hb1f1ab4_1/lib/python3.10/site-packages/torch/include/ATen/core/UnsafeFromTH.h
/usr/share/doc/info/html/INFOPATH.html

这里因为有多个conda环境所以有多个TH.h

建立软链接，

sudo ln -s /home/yunjie/.local/lib/python3.8/site-packages/torch/include/ATen /usr/include/c++/9/    # 可能要sudo，不加可能权限不够

2. 在.bashrc中添加环境变量

在~/.bashrc中为c++添加环境变量，具体语句如下：

export CPLUS_INCLUDE_PATH=/home/yunjie/.local/lib/python3.8/site-packages/torch/include/:$CPLUS_INCLUDE_PATH     # 改为自己的位置就行

添加完成之后，好像一定要重启，不然无法导入。

测试

直接上代码，如果编译能通过则没有问题。

#include <TH/TH.h>      // pytorch的头文件
#include <math.h>
#include <iostream>

int main(int argc, char const *argv[])
{
    std::cout << "chenjun" << std::endl;
    return 0;
}