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

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*/);
//这个不用指定设备
auto/at::Tensor a = at::zeros_like(tensor);

//获得各个维度的大小
auto/int a = tensor.size(dim);

//将tensor复制到GPU上
tensor.to(at::kCUDA)

下面是很重要的把tensor的数据类型和kernel函数进行统一的接口。（代码来自博客
https://blog.csdn.net/u011509971/article/details/104064178，由于沟通具有很大延迟，未经允许就进行了使用，侵权即删。）

AT_DISPATCH_FLOATING_TYPES(input.type(),"ConvBackward", ([&]{
            const scalar_t *bottom_data = input.data<scalar_t>();
            const scalar_t *depthwise_weight = weight.data<scalar_t>();
            const scalar_t *top_diff = output_grad.data<scalar_t>();
            scalar_t *depthwise_weight_diff = weight_diff.data<scalar_t>();
            scalar_t *depthwise_bottom_diff = bottom_diff.data<scalar_t>();

            if (bias_term_){
                scalar_t *depthwise_bias_diff = bias_diff.data<scalar_t>();
                ConvBackwardBias<scalar_t><<<count_bias, THREADS_PER_BLOCK>>>(count_bias, top_diff, batch_size,
                    channels, input_height, input_width, conved_height, conved_width, kernel_h, kernel_w, stride_h,
                    stride_w, pad_h, pad_w, depthwise_bias_diff);
            }

            ConvBackwardWeight<scalar_t><<<count_weight, THREADS_PER_BLOCK>>>(count_weight, top_diff, batch_size,
                channels, input_height, input_width, conved_height, conved_width, kernel_h, kernel_w, stride_h,
                stride_w, pad_h, pad_w, depthwise_weight_diff, bottom_data);

            ConvBackward<scalar_t><<<count_input, THREADS_PER_BLOCK>>>(count_input, top_diff, batch_size,
                channels, input_height, input_width, conved_height, conved_width, kernel_h, kernel_w, stride_h,
                stride_w, pad_h, pad_w, depthwise_bottom_diff, depthwise_weight);
            
        }));

第一行的AT_DISPATCH_FLOATING_TYPES是一个封装的接口，可以替换成AT_DISPATCH_ALL_TYPES。它有三个参数，第一个是tensor的数据类型，第二个是用于显示错误的信息，第三个是个匿名函数，([&]{ })内写cuda的__global__ kernel函数。
input.data() 把input的数据转换成scalar_t类型并且返回一个头指针，该数据是一个一维的连续存储的地址，访问数据的方式和c语言指针使用方法一样。
Pytorch的tensor是按行来存储的，在内存里是一个一维的连续的数组，每次打印tensor的顺序，一行一行展开就是它的存储顺序。tensor.storage()可以打印它的存储顺序，tonsor.data_ptr()返回数据的首地址。

"Tensor.data() is deprecated. Please use Tensor.data_ptr() instead."

而ConvBackwardBias<<>>则是CUDA的语法，<<<,>>>里面的是线程的组织结构和数量，第一个是block的结构，第二个是thread的结构。可以通过blockIdx.x blockId.y threadIdx.x threadIdx.y来访问线程的位置。
具体的CUDA写法和其他内容可以参见其他博客，不想在这里赘述了。

setup.py 和 test.py

这个文件的写法很固定
如下

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

setup(
    name='testname',
    ext_modules=[
        CUDAExtension('test_cuda', [
            './test.cpp',
            './test.cu',
        ])
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

其中testname是当你python setup.py install完成之后，生成的模块的名称，但并不是可以import的名称，test_cuda才是可以import的名称。但是直接import test_cuda是不可以的，需要用一个test.py封装起来，在test.py里最好写完对应的函数testfunction和对应的封装的模块Test。

import math
from torch import nn
from torch.autograd import Function
import torch
import test_cuda
devices=torch.device('cuda')

class testfunction(Function):
    @staticmethod
    def forward(ctx, inputs):
        outputs = test_cuda.forward(inputs)
        ctx.save_for_backward(inputs)
        return outputs

    @staticmethod
    def backward(ctx, grad_output):
        inputs=ctx.saved_variables
        outputs = LSBlock_cuda.backward(grad_output.contiguous(),inputs)
        input_diff= outputs
        return input_diff

class Test(nn.Module):
    def __init__(self):
        super(Test, self).__init__()
        #其他初始化
    def forward(self, inputs):
        return testfunction.apply(inputs)

这里注意testfunction的参数grad_output在传入反传函数时，要加上**.contiguous()**，这让数据的存储是连续的，这很关键。
以后就可以直接导入test.py使用Test模块了

callandrun.py

import sys
sys.path.append(test.py所在目录)
import test  # from test import Test

上述的代码可以让你在任意目录下使用该模块