CUDA实现focal

参考自：mmdetection源码阅读：cuda拓展之focal loss - 知乎

读者需要大致了解CUDA编程及损失函数原理；本文不做详细介绍

CUDA实现加速的写法（套路）

图片来自上述参考文献（侵删），红色文字是我加的注释；

这样说可能还是模模糊糊，下面会详细讲解；

实例讲解：focal_loss cuda实现

注：本节代码内容来自开头参考文章

step1：python端调用（源代码在mmdetection工具包）

注意到调用了 _sigmoid_focal_loss(即sigmoid_focal_loss)，该函数本质上就是对 cuda版本的focal_loss实现的封装，下面介绍；

from mmcv.ops import sigmoid_focal_loss as _sigmoid_focal_loss
# sigmoid_focal_loss其实就是SigmoidFocalLossFunction的forward方法class FocalLoss(nn.module):def forward(self,pred, # tensor(num_total_anchors, num_classes)target, # tensor(num_total_anchors, )):if if torch.cuda.is_available() and pred.is_cuda:loss = _sigmoid_focal_loss(pred.contiguous(), target.contiguous(), gamma,alpha, None, 'none')# 注意要经过contiguous确保内存连续存储！这样在cuda核函数访问连续内存就不会出错return loss

step2：autograd.Function的使用，用来指定前后向计算方法（为什么：我们用cuda定义了一个算子，封装之后需要用autograd让torch知道怎么来做前后向计算）

可以看到继承了Function；使用了ext_module（来自下一步的python-cuda绑定）；使用了Function.apply()

class SigmoidFocalLossFunction(Function):@staticmethoddef forward(ctx,input,target,gamma=2.0,alpha=0.25,weight=None,reduction='mean'):# 存储reduction_dict、gamma等到ctx，以便反传backward调用ctx.reduction_dict = {'none': 0, 'mean': 1, 'sum': 2}assert reduction in ctx.reduction_dict.keys()ctx.gamma = float(gamma)ctx.alpha = float(alpha)ctx.reduction = ctx.reduction_dict[reduction]output = input.new_zeros(input.size()) # 开辟output的空间# 调用真正的cuda拓展：这里ext_module即被用来绑定cuda版本的代码的ext_module.sigmoid_focal_loss_forward(input, target, weight, output, gamma=ctx.gamma, alpha=ctx.alpha)if ctx.reduction == ctx.reduction_dict['mean']:output = output.sum() / input.size(0)elif ctx.reduction == ctx.reduction_dict['sum']:output = output.sum()ctx.save_for_backward(input, target, weight)  # 保存变量供反向计算使用return output@staticmethod@once_differentiabledef backward(ctx, grad_output):input, target, weight = ctx.saved_tensorsgrad_input = input.new_zeros(input.size())# 调用真正的cuda拓展    ext_module.sigmoid_focal_loss_backward(input, target,weight, grad_input, gamma=ctx.gamma, alpha=ctx.alpha)grad_input *= grad_outputif ctx.reduction == ctx.reduction_dict['mean']:grad_input /= input.size(0)return grad_input, None, None, None, None, None# 定义sigmoid_focal_loss为SigmoidFocalLossFunction.apply，apply方法会调用forward
sigmoid_focal_loss = SigmoidFocalLossFunction.apply

step3：使用pybind绑定python-cuda(c++)，以实现调用python的module来调用cuda的函数；

这个过程在c++文件内定义；
TORCH_EXTENSION_NAME会对应到你指定的一个名字，对于上面代码就是 ext_module（也可以叫张三李四，这个一般在setup.py内指定）；
m表示module，def的第一个参数是名字(name)，第二个参数是要绑定的c++/cuda函数，这里就是 ext_module.sigmoid_focal_loss_forward. 绑定到当前c++文件定义的 sigmoid_focal_loss_forward函数；这个函数就是真正涉及到CUDA加速的函数；下面介绍；
def 的其他参数就是相关函数的参数了，不影响此处理解，不做赘述
sigmoid_focal_loss_forward函数将会调用SigmoidFocalLossForwardCUDAKernelLauncher函数；见下文

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {  m.def("sigmoid_focal_loss_forward", &sigmoid_focal_loss_forward,"sigmoid_focal_loss_forward ", py::arg("input"), py::arg("target"),py::arg("weight"), py::arg("output"), py::arg("gamma"),py::arg("alpha"));
}
# 同理sigmoid_focal_loss_backward略

step4：SigmoidFocalLossForwardCUDAKernelLauncher，根据各种情况设置CUDA线程的组织形式以及其他一些基本设置（可以理解为准备好环境），然后调用真正的CUDA kernel进行核心部分计算（这里就是focal_loss的计算了，真正干活的）

Launcher这个名字比较形象，但是读者在其他地方不一定会见到Launcher这个写法，因为这个不是规定；但是肯定会有相关函数实现这部分功能；

以下是c++代码：

获取了相关参数并做一些基本判断；还有设置线程组织形式，就是<<<a, b>>>内的a,b
准备好环境后，就可以调用kernel了；
AT_DISPATCH_FLOATING_TYPES_AND_HALF是个带参数的宏，干的事就是把第一个参数（数据类型）传给后面的scalar_t，然后执行[&]()括号里的匿名函数（就是kernel了）

此外还注意到Tensor提供了.data_ptr<scalar_t>()模板成员函数，将返回tensor的连续存储的首地址，并且转换为scalar_t *的指针类型。我们知道tensor在内存中真正存放的是一维连续数组！tensor(B,C,H,W)在内存中真正存放的是长为B*C*H*W的连续数组；data_ptr()返回的就是这个连续数组的首地址；

void SigmoidFocalLossForwardCUDAKernelLauncher(Tensor input, Tensor target,Tensor weight, Tensor output,const float gamma,const float alpha) {
// input为tensor(num_total_anchors, num_classes)
// output为tensor(num_total_anchors, num_classes)
// target为tensor(num_total_anchors, )，0 ~ num_class-1表示正样本对应的类别，num_class值表示负样本和忽略样本int output_size = output.numel(); //等于num_total_anchors*num_classesint num_classes = input.size(1);AT_ASSERTM(target.max().item<long>() <= (long)num_classes,"target label should smaller or equal than num classes");at::cuda::CUDAGuard device_guard(input.device());cudaStream_t stream = at::cuda::getCurrentCUDAStream();AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "sigmoid_focal_loss_forward_cuda_kernel", [&] {sigmoid_focal_loss_forward_cuda_kernel<scalar_t> //sigmoid_focal_loss_forward_cuda_kernel是cuda核函数，被定义成模板函数，通过<scalar_t>确定数据类型<<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(output_size, input.data_ptr<scalar_t>(),target.data_ptr<int64_t>(), weight.data_ptr<scalar_t>(),output.data_ptr<scalar_t>(), gamma, alpha, num_classes);});AT_CUDA_CHECK(cudaGetLastError());
}//backward略// 以下是CUDA计算需要设定的线程组织形式，不熟悉的请查阅CUDA编程
#define THREADS_PER_BLOCK 1024、128、512inline int GET_BLOCKS(const int N) {int optimal_block_num = (N + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;int max_block_num = 65000;return min(optimal_block_num, max_block_num);
}
// 获得开辟的线程块的x维，使得blockDim.x * gridDim.x要等于N

step5：kernel实现

这就是真正干活的了；这时候就得理解focal_loss的计算方式，才知道怎么写了；（其实上面在定义<<<a,b>>>里面的a、b参数时，就应该考虑好这部分的写法）
先推导一下focal_loss吧，过程如下
接下来就是实现了，kernel的内容代表了每个线程要干啥事；我们这里要计算loss，而且是分类的loss；具体来说，输入包括
- pred，是一个tensor，shape为[num_total_anchors, num_classes]；
- 以及target，也是一个tensor，shape为[num_total_anchors]；
- 可能还有weight，表示每个类别的损失权重（有些类别不重要，就给小权重，这样即便预测损失大，最后因为权重小对平均损失影响不大）
通常实现这种多分类loss，就是softmax；mmdetection使用的是sigmoid+focal_loss；就是利用sigmoid对num_classes里面所有类别计算损失，结合focal_loss（就是正负类加不一样的权重），最后做综合（比如取平均得到平均损失）；

有了上述思路，我们就知道如何设计CUDA的多线程计算（就是设计kernel了）；

既然要对每个anchor的num_classes的每个类别计算损失（此时总共要计算N=num_total_anchors*num_classes这个多个损失），那么就让它们都并行计算好了；
所以我们需要 num_total_anchors*num_classes 这么多线程；这就是上面 <<<a,b>>>里a的GET_BLOCK参数设置为output_size= num_total_anchors*num_classes的原因；当然可能CUDA并没有这么多可同时用的线程，没关系，再设置个CUDA范围内的总线程数呗，剩下的工作量让这些线程循环执行；这就是上面GET_BLOCK干的事，请看下文
CUDA要求设置线程块大小（一般叫THREADS_PER_BLOCK）以及线程块数量（一般叫BLOCKS or BLOCK_NUM or BLOCKS_PER_GRID）；通常两者乘积尽量稍大于总任务数（比如上面需要计算的损失个数N）；但是CUDA没那么多线程，所以也得控制一下两个常量的大小，比如上面GET_BLOCK里最后设置max控制线程块数量了；
现在有THREADS_PER_BLOCK*BLOCKS_PER_GRID这么多线程了，可以理解为它们会在CUDA内部并行执行，每个线程计算一个损失，剩下的损失由前面的线程循环执行；

根据上面几步的介绍以及上面focal_loss前后向计算的推导，我们可以写出如下kernel，几个注意点：

CUDA_1D_KERNEL_LOOP(i, n)是个带参数宏，用来定义循环，就是上面说的线程不够循环凑
- i是CUDA自动计算的当前线程的总index，n则是上面计算的总任务数（所有需要计算的损失的数量，即output_size）
所以index用来定位当前线程应该计算哪个损失；注意前面提到了，多维的tensor在内存中实际上是一维形式存在的，所以kernel的参数是指针形式的；此时用input[index]就找到了相应的位置了；然后可以计算当前是第几个anchor(文中n)和该anchor的第几类的损失了；
接下来的计算就是根据上面的focal_loss的前后向推导公式了；整体算是比较简单的
- 注意后向计算（backward函数），并行线程数可以与前向计算forward不一样，但是计算方式得与forward协调一致（比如对每个位置计算loss的话就要对每个位置计算grad）

代码如下

#define CUDA_1D_KERNEL_LOOP(i, n)                            \for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \i += blockDim.x * gridDim.x)
// blockDim.x * gridDim.x就是当前开辟的线程总数template <typename T>
__global__ void sigmoid_focal_loss_forward_cuda_kernel(const int nthreads, const T* input, const int64_t* target, const T* weight,T* output, const T gamma, const T alpha, const int num_classes) {
// nthreads就是outputsize，等于num_total_anchors*num_classes
// const T*就是tensor的连续内存首地址
// input的连续内存长度为num_total_anchors*num_classes，target的连续内存长度为num_total_anchors。CUDA_1D_KERNEL_LOOP(index, nthreads) {// index等于blockIdx.x * blockDim.x + threadIdx.x，即线程索引// 因为index就是对应tensor(num_total_anchors,num_classes)的一个元素int n = index / num_classes; // 所以n就是该元素对应的anchorint c = index % num_classes; // 所以c就是该元素对应的classint64_t t = target[n]; // 获得anchor n 的target labelT flag_p = (t == c); // 表示正样本T flag_n = (t != c); // 表示负样本// p = sigmoid(x) = 1. / 1. + expf(-x)T p = (T)1. / ((T)1. + expf(-input[index]));// (1 - p)**gamma * log(p) 正样本的focal loss权重T term_p = pow(((T)1. - p), gamma) * log(max(p, (T)FLT_MIN));// p**gamma * log(1 - p) 负样本的focal loss权重T term_n = pow(p, gamma) * log(max((T)1. - p, (T)FLT_MIN));output[index] = (T)0.; //计算结果放到output tensor中output[index] += -flag_p * alpha * term_p;output[index] += -flag_n * ((T)1. - alpha) * term_n;if (weight != NULL) {output[index] *= weight[t];}}
}//同理，反向传播
template <typename T>
__global__ void sigmoid_focal_loss_backward_cuda_kernel(const int nthreads, const T* input, const int64_t* target, const T* weight,T* grad_input, const T gamma, const T alpha, const int num_classes) {CUDA_1D_KERNEL_LOOP(index, nthreads) {int n = index / num_classes;int c = index % num_classes;int64_t t = target[n];T flag_p = (t == c);T flag_n = (t != c);// p = sigmoid(x) = 1. / 1. + expf(-x)T p = (T)1. / ((T)1. + exp(-input[index]));// (1 - p)**gamma * (1 - p - gamma*p*log(p))T term_p = pow(((T)1. - p), gamma) *((T)1. - p - (gamma * p * log(max(p, (T)FLT_MIN))));// p**gamma * (gamma * (1 - p) * log(1 - p) - p)T term_n = pow(p, gamma) *(gamma * ((T)1. - p) * log(max((T)1. - p, (T)FLT_MIN)) - p);grad_input[index] = (T)0.;grad_input[index] += -flag_p * alpha * term_p;grad_input[index] += -flag_n * ((T)1. - alpha) * term_n;if (weight != NULL) {grad_input[index] *= weight[t];}}
}

总结

以上介绍了mmdetection中的focal_loss算子的CUDA实现；
算子实现有比较固定的流程，由顶向下分别是：
1. Python调用autograd Function
2. 上述autograd定义前后向算法
3. 上述前后向算法引用Python用来绑定c++函数的module
4. 上述module有前后向函数映射到c++的前后向函数
5. c++的前后向函数准备环境（最重要的是线程组织形式）并调用CUDA kernel
6. kernel定义CUDA如何实现具体算子（如上面的focal_loss）
难点在于如何实现最后kernel；具体的任务要先能分解成多个线程能并行同时做的，才好实现kernel；
TODO：其他更复杂算子的CUDA实现

CUDA实现focal_loss相关推荐

U-net 训练自己的数据集
文章目录一.网络准备二.数据集准备 2.1 JSON标签数据转png 2.1.1 安装LABELME 2.1.2 制作标签图片 2.1.3 将json转换格式三. win10+pycharm训练 ...
Keras自定义损失函数出现：ValueError: Unknown loss function: focal_loss
Keras自定义损失函数出现:ValueError: Unknown loss function: focal_loss 1.软件环境 2.问题描述 3.解决方法 4.结果预览 1.软件环境 Wind ...
torch.backends.cudnn.deterministic 使用cuda保证每次结果一样
为什么使用相同的网络结构,跑出来的效果完全不同,用的学习率,迭代次数,batch size 都是一样?固定随机数种子是非常重要的.但是如果你使用的是PyTorch等框架,还要看一下框架的种子是否固定了 ...
RuntimeError: Expected object of backend CUDA but got backend CPU for argument
RuntimeError: Expected object of backend CUDA but got backend CPU for argument #4 'mat1' 原因:变量没有加cud ...
RuntimeError: Expected object of device type cuda but got device type cpu for argument pytorch数据位置
RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'target' i ...
cuda安装教cudnn安装显卡问题
20220101 https://www.cnblogs.com/zucc-31701019/p/14665855.html CUDA 安装报错 could not create file " ...
在Cuda上部署量化模型
在Cuda上部署量化模型介绍TVM自动量化.自动量化是TVM中的一种量化方式.将在ImageNet上导入一个GluonCV预先训练的模型到Relay,量化Relay模型,然后执行推理. import ...
CUDA上深度学习模型量化的自动化优化
CUDA上深度学习模型量化的自动化优化深度学习已成功应用于各种任务.在诸如自动驾驶汽车推理之类的实时场景中,模型的推理速度至关重要.网络量化是加速深度学习模型的有效方法.在量化模型中,数据和模型参数 ...
CUDA 11功能清单
CUDA 11功能清单基于NVIDIA Ampere GPU架构的新型NVIDIA A100 GPU在加速计算方面实现了最大的飞跃.A100 GPU具有革命性的硬件功能,CUDA 11与A100一起 ...
CUDA上的量化深度学习模型的自动化优化
CUDA上的量化深度学习模型的自动化优化深度学习已成功应用于各种任务.在诸如自动驾驶汽车推理之类的实时场景中,模型的推理速度至关重要.网络量化是加速深度学习模型的有效方法.在量化模型中,数据和模型参 ...

CUDA实现focal_loss

CUDA实现加速的写法（套路）

实例讲解：focal_loss cuda实现

总结

CUDA实现focal_loss相关推荐

最新文章

热门文章