CUDA实现focal_loss

参考自: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):
    @staticmethod
    def 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_differentiable
    def backward(ctx, grad_output):
        input, target, weight = ctx.saved_tensors

        grad_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_output
        if 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
    • 准备好环境后,就可以调用kernel了;
    • AT_DISPATCH_FLOATING_TYPES_AND_HALF是个带参数的宏,干的事就是把第一个参数(数据类型)传给后面的scalar_t,然后执行[&]()括号里的匿名函数(就是kernel了)
    • 此外还注意到Tensor提供了.data_ptr()模板成员函数,将返回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_classes
        int num_classes = input.size(1);
        AT_ASSERTM(target.max().item() <= (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 
             //sigmoid_focal_loss_forward_cuda_kernel是cuda核函数,被定义成模板函数,通过确定数据类型
                  <<>>(
                      output_size, input.data_ptr(),
                      target.data_ptr(), weight.data_ptr(),
                      output.data_ptr(), gamma, alpha, num_classes);
            });
      
        AT_CUDA_CHECK(cudaGetLastError());
      }
      
      //backward略
      
      // 以下是CUDA计算需要设定的线程组织形式,不熟悉的请查阅CUDA编程
      #define THREADS_PER_BLOCK 1024、128、512
      
      inline 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参数时,就应该考虑好这部分的写法)
  • 先推导一下focal_loss吧,过程如下CUDA实现focal_loss_第1张图片
  •  接下来就是实现了,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设置为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 
        __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就是该元素对应的anchor
            int c = index % num_classes; // 所以c就是该元素对应的class
        
            int64_t t = target[n]; // 获得anchor n 的target label
            T 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 
        __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实现;
  • 算子实现有比较固定的流程
  • 难点在于如何实现最后kernel;具体的任务要先能分解成 多个线程能并行同时做的,才好实现kernel;
  • TODO:其他更复杂算子的CUDA实现

你可能感兴趣的:(技术问题,#,CUDA,#,深度学习,深度学习,计算机视觉,CUDA)