撰文｜郑泽康

InsightFace模型里大量应用了PReLU激活函数，而PReLU的工作模式有两种：

  1. PReLU(1)，此时权重alpha的形态为(1, )，等价于一个Elementwise操作。
  2. PReLU(channels)，此时权重alpha的形态为(channels, )，和输出特色(N, C, H, W)中C的大小是对应的。此时PReLU等价于一个Binary Broadcast操作。

InsightFace模型里的PReLU工作模式是第二种，之前曾经介绍过CUDA Elementwise操作优化，而在Broadcast情景下也存在肯定的优化机会。

1

奢侈实现

一个奢侈实现的思维就是在循环外部，依据以后元素的索引，推算出该元素对应须要应用的alpha权重的索引。而后判断以后元素x是否大于0，若大于0则返回x，小于0则返回alpha*x。对应代码如下：

template<typename T>__global__ void PReluForwardGpu(const int32_t elem_cnt, const int32_t alpha_size,                                const int32_t inner_size, const T* x, const T* alpha, T* y) {  CUDA_1D_KERNEL_LOOP(i, elem_cnt) {    const T x_i = x[i];    const T alpha_i = alpha[(i / inner_size) % alpha_size];    y[i] = x_i > 0 ? x_i : x_i * alpha_i;  }}

其中：

  • inner_size示意的是通道维前面维度乘积，以NCHW格局为例，inner_size=H*W
  • alpha_size示意通道维大小

在CUDA中，整数除法的计算代价是比拟低廉的（https://docs.nvidia.com/cuda/...）对于计算指令耗时这一章中有提到：

Integer division and modulo operation are costly as they compile to up to 20 instructions.

整数除法，取余操作会被编译成多达20条指令。而咱们这里计算alpha的索引的时候，别离用到一次除法，一次取余，占整个Kernel的次要计算量，上面咱们将用向量化的思路来进步读写带宽的同时，缩小整数除法，取余的计算次数。

2

Pack向量化优化

咱们思考一个比较简单的例子，输出为(1, 2, 4, 4)，对应PReLU(2)

显然，输出在hw维上是间断的，在 inner_size 满足被pack整除的条件下，一个pack内的元素利用到的是同一个alpha权重**。参见下图：

这样咱们就能以向量化模式去解决元素，以晋升读写带宽。并且每一个pack外部只须要计算一次，向量化解决相比逐元素计算能节俭不小计算量。对应代码如下：

template<typename T, typename IndexType, int pack_size>__global__ void PReluForwardMultiAlphaGpu(const IndexType elem_cnt, const IndexType alpha_size,                                          const IndexType inner_size, const T* x, const T* alpha, T* y) {  int32_t global_thread_id = blockIdx.x * blockDim.x + threadIdx.x;  using LoadType = cuda::elementwise::PackType<T, pack_size>;  using LoadPack = cuda::elementwise::Pack<T, pack_size>;  T zero_val = static_cast<T>(0);  for (int64_t linear_index = global_thread_id * pack_size; linear_index < elem_cnt;       linear_index += gridDim.x * blockDim.x * pack_size) {    // 计算以后Pack所应用到Alpha的索引        IndexType alpha_idx = (linear_index/inner_size%alpha_size);    const LoadType* x_load = reinterpret_cast<const LoadType*>(x + linear_index);    // 以向量化的模式加载输出x    LoadPack x_vec;    x_vec.storage = *x_load;    LoadPack y_vec;    // 循环展开，一一解决Pack内的元素#pragma unroll    for (int i = 0; i < pack_size; i++) {      y_vec.elem[i] = x_vec.elem[i] > zero_val ? x_vec.elem[i] : x_vec.elem[i] * alpha[alpha_idx];    }    // 以向量化的模式存储输入y    *(reinterpret_cast<LoadType*>(y + linear_index)) = y_vec.storage;  }}

咱们在Nsight Compute内简略比拟下优化前后的后果，测试数据为(96, 64, 112, 112)，机器为A100-40GB。蓝色一栏是应用向量化优化过的kernel，而绿色一栏是奢侈实现的kernel。能够看到，通过优化后，咱们计算占比升高20%-30%，吞吐晋升了30+%。优化后的kernel带宽能达到1350GB/s，曾经很靠近A100上的实践带宽1555GB/s。

当然也不是所有形态都反对向量化操作，当inner_size无奈被对应的pack_size 整除时，只能退回到奢侈实现上。

3

基准测试

在A100-40GB测试机器上，咱们对Insightface波及到的Tensor形态，与PyTorch实现进行比拟，测试数据如下：

通过优化PReLU的OneFlow，在大部分状况下均有比PyTorch靠近2倍的当先劣势，在最初一种状况因为形态较为非凡，无奈利用向量化的优化，所以体现与PyTorch持平。

欢送下载体验 OneFlow v0.7.0 最新版本：
https://github.com/Oneflow-In...