在深度学习任务中,如果使用PyTorch实现神经网络中的一些功能模块,其运行速度是比较慢的;为了能够加快这些模块的运行速度,通常会使用CUDA或者Cython完成这些功能模块的逻辑编写,并编译成动态链接库,之后在Python文件中实现调用,并封装成相应的功能模块,实现网络的前向和反向传播,从而大大提高网络整体的执行速度。
本文主要是想梳理一下实现上述功能的整体流程。
- 如何使用CUDA完成某一功能模块的编写;
- 编写好的CUDA扩展如何能够和Python进行连接/交互 (后面会提到,这一连接主要是靠C++文件);
- 如何将写好的CUDA代码,C++代码进行混合编译;
- 编写好的CUDA扩展,如何被Python封装成网络的一个模块,并实现网络的前向和反向传播;
1.文件整体结构
想要实现CUDA和C++的混合编译,并能够被Python调用,执行相应的功能主要需要4个文件,如下所示。主要参考Yulong:扩展Pytorch:利用CUDA实现算子(二)
(一) CPP文件
在CPP文件中主要完成4个功能。
1)对要调用的CUDA函数进行声明,这里提到的CUDA函数的定义在cu文件中;
2)完成一些宏定义,这些宏定义主要作用是为了进行数据的形状检查或张量检查,防止运行过程中出现问题;
3)定义Python调用扩展功能的函数接口,也就是CPP文件中定义的C++函数;
4)PYBIND11_MODULE(TORCH_EXTENSION_NAME)函数封装Python调用扩展函数的C++接口。
(二)CU文件
在CU文件中一般包括2个文件。
1)对CPP文件中的声明的CUDA函数进行定义;在该函数中一般会调用真正实现模块逻辑功能的kernel函数;
2)定义实现模块真正逻辑功能的kernel函数。
注:在实际项目中,一般这种编译的扩展文件都会被集中放在一个名为src的文件夹中,在这个src文件夹中通常根据不同的功能包含很多的xxx文件夹,比如本文扩展想要实现sigmoid()函数的运算,那么CPP文件和CUDA文件就会放在sigmoid文件夹下。
那如何实现各个功能模块的编译呢,就是靠下面的setup.py文件啦~
(三)setup.py文件
实现CPP文件和CUDA文件的混合编译。
(四)test.py文件
在Python代码中实现对扩展的调用,并将相应功能封装成模块,定义相应的forward()函数和backward()函数,实现网络的前向和反向传播。
下面以实现一个非常简单的sigmoid()运算为例,介绍下如何编写上述的4个文件。
2. CPP文件
由于上面已经整体介绍过了,这里主要说一下PYBIND11_MODULE函数中m.def()3个变量是什么含义:
1)第一个变量"forward":实际就是在Python文件中调用CPP文件中定义的C++函数的接口;
2)第二个变量&sigmoid_forward:这个位置的名称一定要和C++定义的函数名称一致,并前面加上引用&,这样在Python文件中调用forward函数接口,才会调用CPP文件中定义的sigmoid_forward()函数,从而完成后续一系列函数的调用;
3)第三个变量"sigmoid forward (CUDA)":相当于注释,方便别人阅读或自己知道这个函数的功能。
注:在编译这个CPP代码时,使用torch::Tensor或者at::Tensor命名空间都是可以通过编译的。而且即使上面没有include Aten.h相关的库,直接使用at::Tensor也没有问题 ,希望有大佬可以解答一下,为啥行呀? 😃。
#include<torch/torch.h>
#include<vector>
// declare cuda forward function
torch::Tensor sigmoid_cuda_forward(torch::Tensor input);
// declare cuda backward function
torch::Tensor sigmoid_cuda_backward(torch::Tensor output);
// macro definition
#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), "data must be a CUDA Tensor.")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), "data must be contiguous.")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
// define c++ forward function
torch::Tensor sigmoid_forward(torch::Tensor input){
CHECK_INPUT(input);
return sigmoid_cuda_forward(input);
}
// define c++ backward function
torch::Tensor sigmoid_backward(torch::Tensor output){
CHECK_INPUT(output);
return sigmoid_cuda_backward(output);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
m.def("forward", &sigmoid_forward, "sigmoid forward (CUDA)");
m.def("backward", &sigmoid_backward, "sigmoid backward (CUDA)");
}
3. CUDA文件
上述也提到过,CUDA文件中主要包含两种函数,一种是呼应CPP里声明的CUDA函数,另一种是真正实现功能的kernel函数。
在本文的例子中,sigmoid_cuda_forward()以及sigmoid_cuda_backward()函数就是呼应CPP的函数,并在这些函数中调用真正实现模块逻辑功能的kernel函数,也就是sigmoid_cuda_forward_kernel()和sigmoid_cuda_backward_kernel()函数;
其中AT_DISPATCH_FLOATING_TYPES()函数的功能,参考 Pytorch拓展C++和CUDA:讲解和细节-爱代码爱编程
注:在CUDA文件中,就只能使用at::Tensor而不能使用torch::Tensor,一旦使用torch::Tensor就会报错,不知道什么原因,希望有大佬可以解答下我的疑惑 😃。
#include<ATen/ATen.h>
#include<cuda.h>
#include<cuda_runtime.h>
#include<vector>
template <typename scalar_t>
__device__ scalar_t sigmoid(scalar_t z){
return 1.0 / (1.0 + exp(-z));
}
template <typename scalar_t>
__device__ scalar_t d_sigmoid(scalar_t z){
return (1.0 - z) * z;
}
template <typename scalar_t>
__global__ void sigmoid_cuda_forward_kernel(const scalar_t * __restrict__ input, scalar_t * __restrict__ output){
const int index = blockIdx.x * blockDim.x + blockIdx.y;
output[index] = sigmoid(input[index]);
}
template <typename scalar_t>
__global__ void sigmoid_cuda_backward_kernel(const scalar_t* __restrict__ output,
scalar_t* __restrict__ new_grad_output){
const int index = blockIdx.x * blockDim.x + blockIdx.y;
new_grad_output[index] = d_sigmoid(output[index]);
}
// only using at::Tensor in .cu file
// not using torch::Tensor
at::Tensor sigmoid_cuda_forward(at::Tensor input){
auto output = at::zeros_like(input);
dim3 blocks(input.size(0), input.size(1));
int threads = 1;
AT_DISPATCH_FLOATING_TYPES(input.type(), "error in sigmoid_cuda_forward", ([&]
{sigmoid_cuda_forward_kernel<scalar_t> <<<blocks, threads>>> (input.data<scalar_t>(), output.data<scalar_t>());
}));
return output;
}
at::Tensor sigmoid_cuda_backward(at::Tensor output){
auto new_grad_output = at::zeros_like(output);
dim3 blocks(output.size(0), output.size(1));
int threads = 1;
AT_DISPATCH_FLOATING_TYPES(output.type(), "error in sigmoid_cuda_backward", ([&]{
sigmoid_cuda_backward_kernel<scalar_t> <<<blocks, threads>>> (output.data<scalar_t>(),
new_grad_output.data<scalar_t>());}));
return new_grad_output;
}
4. setup.py文件
该文件主要就是完成上述代码的混合编译,了解的不是很多,就不多介绍了。
# setup()函数中变量的含义
name: 在python中,import该模块的名称
sources: 源代码文件的名称
laungage: 默认为c,可以改成c++
include_dirs: 传递给gcc或者g++编译器,include的头文件目录
library_dirs: 传递给gcc或者g++编译器,指定链接文件的目录
libraries: 传递给gcc或者g++编译器,指定链接的文件
extra_compile_args: 额外的gcc或者g++编译参数
extra_links_args: 额外的传给g++或者gcc编译器的链接参数
define_macros: 定义的宏
undefine_macros: 取消宏
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='_CUDA',
ext_modules=[
CUDAExtension('sigmoid_cuda', [
'src/sigmoid_cuda.cpp',
'src/sigmoid_cuda_kernel.cu',
]),
],
cmdclass={
'build_ext': BuildExtension
})
5. test.py文件
该文件就是将扩展功能封装到一个模块中,并在该模块中定义前向传播和反向传播函数。之所以要定义反向传播函数是因为模块是自己用CUDA写的,无法使用PyTorch自带的反向求导,所以必须要自己实现反向传播的逻辑。具体的实现思路可以参考 科技猛兽:PyTorch 74.自定义操作torch.autograd.Function
这里主要呼应一下上面提到的m.def()函数第一个参数的"用武之地"。
from torch.nn import Module
from torch.autograd import Function
import torch
import sigmoid_cuda
class SigmoidFunction(Function):
@staticmethod
def forward(ctx,
input):
# 在cpp文件中定义为"forward",
# 在这个位置就用.forward()实现对CPP文件中的C++函数进行调用
output = sigmoid_cuda.forward(input)
print(f'forward result: {output}')
ctx.save_for_backward(output)
return output
@staticmethod
def backward(ctx,
grad_output):
output = ctx.saved_tensors
output = output[0]
grad_sigmoid = sigmoid_cuda.backward(output.contiguous())
grad_result = grad_sigmoid * grad_output
print(f'backward result: {grad_result}')
return grad_result
class Dense(Module):
def __init__(self):
super(Dense, self).__init__()
def forward(self, input):
return SigmoidFunction.apply(input)
if __name__ == '__main__':
a = torch.tensor([[1.], [3.]], dtype=torch.float32, requires_grad=True).cuda()
print(f'input data: {a}')
m = Dense()
result = m(a)
sum_ = result.sum()
sum_.backward()
6. 代码的具体实现链接
理论说了这么多,毕竟“纸上得来终觉浅,绝知此事要躬行”,我将上面的所有代码整理到了repo中,这个repo中还有我收集的其他扩展功能,欢迎大家PR、Star~。
https://github.com/HsLOL/ExtensionOPs/tree/master/demo
github.com/HsLOL/ExtensionOPs/tree/master/demo
编辑于 2022-05-16 19:07