Pytorch拓展C++和CUDA：讲解和细节

最新推荐文章于 2024-07-16 13:27:06 发布

原创

最新推荐文章于 2024-07-16 13:27:06 发布 · 4.3k 阅读

21 ·

CC 4.0 BY-SA版权

文章标签：

#python #cuda #c++

Pytorch拓展C++ 以及 CUDA

本文不给出一个完全可以运行的完整实例，更重要的讲解如何写以及一些细节。其他博客讲解的官网上的lltm的实例这里并没有，并且最好看过lltm的代码和博客，这样效果更佳。各位看官根据需求选择看或不看。

文件结构和运行方式

#method 1
code/
    test.cpp
    test.cu
    setup.py
    test.py
    callandrun.py
#method 2
code/
    test.cpp
    test.cu
setup.py
test.py
callandrun.py

首先新建一个目录（被很多文章误导，其实这个目录可以随便在一个地方建），目录里包含四个文件
test.cpp文件是起到了连接cuda代码和python代码的作用，cuda代码起到了核心的并行加速的作用，使用setup.py来将其编译生成一个python可以直接import的模块，但是它不能直接见人，需要使用test.py将这个模块封装起来。最后在callandrun.py里调用运行。

test.cpp 和test.cu

这两个文件是核心，其中用到了cuda编程基础，pytorch的C++接口，关键的就是ATen库
ATen库里关于at::Tensor的属性和用法，在这个Tensor类属性链接里，很重要的一个就是at命名空间的用法，在这里 at命名空间
这里给出.slice的用法，其他属性都可以在文档里找到

Tensor slice(int64_t dim = 0, int64_t start = 0, int64_t end = 9223372036854775807, int64_t step = 1) const

torch::Tensor 和at::Tensor其实是一样的，下面给出几个用到的语法规则

//连接两个tensor生成一个新的tensor，重点是两者必须在同一个设备上，同为GPU或者同为CPU
auto/at::Tensor a = at::cat({
   
   tensor1,tensor2},dim);

//生成0 tensor
//根据我的经验 设备是需要指定的 其他人的博客没有写，可能和环境或者cuda的版本或者我写代码的方式有关系
auto/at::Tensor a = at::zeros({
   
   shape1,shape2,...}/*,at::kCUDA*/