卷积实现—im2col+gemm

最新推荐文章于 2024-07-15 22:09:20 发布

two_apples

最新推荐文章于 2024-07-15 22:09:20 发布

阅读量1.5k

点赞数 3

分类专栏：推理框架文章标签： c++ 卷积计算 onnxruntime

本文链接：https://blog.csdn.net/two_apples/article/details/131383786

版权

推理框架专栏收录该内容

1 篇文章 0 订阅

订阅专栏

普通卷积

看卷积的实现，先看其普通的计算方式：滑窗计算和其计算shape大小的公式，以及各个卷积特性对其计算的影响，比如：stride，group，dilation，pad等。 $H_{out}=(H_{in}-k_{h}+p_t+p_b)/stride_{h}+1$ $W_{out}=(W_{in}-k_{w}+p_l+p_r)/stride_{w}+1$ 上面公式以2D卷积为例，输出特征图的形状计算公式。

Group

    假如group不是1，则是分组卷积，分组卷积与普通卷积的区别就是利用分组，将卷积计算的范围更小化，从而降低计算量。分组卷积将一个正常卷积划分为多个更小范围的卷积的拼接。
在这里插入图片描述
    如图这是一个组数为2的分组卷积示意图，其中对输入的特征图通道与滤波器的通道都进行了分组，设置原输入特征图大小为 $NC_{in}HW$ ,滤波器大小为 $C_{out}C_{in}K_{h}K_{w}$ ,输出特征图的输出大小为 $NC_{out}H_{out}W_{out}$ ,假如加入分组卷积系数g后，每个计算卷积参数就会变成 $\frac{C_{out}}{g}\frac{C_{in}}{g}K_{h}K_{w}$ ,得到输出
    矩阵大小为 $N\frac{C_{out}}{g} H_{out}W_{out}$ 。然后由g组的卷积计算结果concat 成为一个正常的卷积输出大小。这样计算量会小很多，但是具体的计算量的降低我们这里就不再多说，这里只是对分组卷积的一个简单描述。

dilation

    感觉在卷积中比较难理解的系数就是一个分组卷积系数，另外一个是空洞卷积系数，空洞卷积的主要目的是为了扩大卷积计算过程中的感受野，相当于将一个卷积扩张成一个更大的卷积，如下图：
在这里插入图片描述
    而关于空洞系数为有效位置之间的关系如下图：

    其中扩张后像素对应位置计算相比还比较重要，因为在卷积计算滑动过程中有效位置的像素点会参与计算，但是无效位置的像素点则会不参与计算。正常卷积进行空洞扩张后，相当于是一个膨胀的卷积，3x3膨胀为5x5的卷积或者7x7的卷积，计算公式如下： $k^{'}=k+(k-1)*(d-1)$     化简后可以表示为： $k^{'}=kd-d+1$     其中相对位置计算也比较重要，33的卷积对应的序号为（i,j），d=2时候，其中像素的对应位置为(i2,j2),当d为其他值时，也可以表示成（id,j*d）。

padding

Padding 其实平时还是比较熟悉的，就是在原特征图周边进行补0操作，才使得卷积操作后得到的特征图是自己想要的shape大小，根据pad求对应的输出特征大小，这里说一下一些关于padding的细节。
padding的出现其实是为了解决一些卷积操作中会存在的一些问题：1.卷积操作会使得输出的图像越来越小，图像变小失去很多有用信息。2.特征图中心区域数据会参与多次卷积计算，而边缘数据信息会被丢失。

Valid卷积

Valid卷积意味着不填充，计算公式直接由p=0时的计算公式1得到。

Same卷积

same卷积主要目的是为了保持卷积输出特征图与输入特征图shape保持一致。same填充一般会出现在stride=1的卷积中，same填充方式会引入大量的边缘0（stride>1），而大量的边缘填充在提取数据特征时候没有什么好处。
计算公式: $H_{out}=ceil(H_{in}/stride_{h})$ $W_{out}=ceil(W_{out}/stride_{w})$ 且在same模式下有可能会出现不对称填充的情况。当滤波器为奇数的时候，会进行自然填充，但是当滤波器为偶数的时候会出现非对称填充，这种情况其实比较麻烦，因为不同的框架中处理非对称填充的方式也不一样，caffe中只支持对称填充，TF中优先填充右下，然后填左上，mxnet的填充优先级和TF相反，先左上再右下。在网络设计中，还是建议大家把padding的形式写明，是same还是valid，因为不同的地方默认的的卷积补充方式不一样，比如tf2与tf2trt接口。tf2默认的padding方式是same，而tf2trt里面的padding默认方式是valid。

Im2col+gemm

其实关于im2col+gemm的实现方式在[1][2]中也阐述的很清楚，但是上面其实更多的是阐述这样为啥可以这样做和在cpu端的设计方式，cpu端和GPU端的设计原理虽然一样，但是GPU其实更有挑战性一点。

原理

以2D卷积过程为例，im2col+gemm的卷积实现其实就是将普通卷积滑窗式计算方案中的滑窗过程展开，将滑窗过程中遍历的特征图像素值保存在一个临时存储空间中的矩阵中，然后再将这个临时存储空间的矩阵与权重进行gemm计算。
Im2col变换方式其实很简单：过程如下图：

滑窗过程整个进行完可以得到：
在这里插入图片描述
im2col这个展开只是针对输入特征图像的，并不这针对权重weight，权重weight不需要做什么操作，权重weight按照NHWC的数据格式存放在内存空间中，依次取出来进行gemm想乘就可以。
另外im2col+gemm就可以完成卷积计算，gemm相乘之后得到的数值就是卷积计算会输出的数值，不需要col2im这个逆过程去转换，这一点[1]说的不太对，可以自己画矩阵直接从数学公式的角度去理解。

设计过程

im2col

    上面是对可行性的原理进行阐述，在im2col+gemm的计算方案设计过程中，对存储逻辑的理解比较细节，大家可以浏览[3]中对于im2col过程描述，感觉描述的还比较好，在设计过程中，我们需要清楚我们需要将原像素对应到临时存储位置（我们这里叫他col_buffer）的哪个位置，暂时以单通道特诊图为例子，在原特征图位置中一个滤波器里面kk的元素做了im2col的展开后会是怎样的存储结构。然后再对应多通道又是怎么对应存储的。一个输入特征图: $N,C_{in},H_{in},W_{in}]$ ，当N取1，并且引入分组卷积后的输入特征图可以看成 $[\frac{C_{in}}{g},H_{in},W_{in}]$ ,im2col过程后会展开成一个二维矩阵： $[\frac{C_{in}}{g}*k*k,H_{out}*W_{out}]$ ,由于存储连续，可以看成是一个三维 $[\frac{C_{in}}{g}*k*k,H_{out},W_{out}]$ ,或者是一个四维 $[\frac{C_{in}}{g}，k*k,H_{out},W_{out}]$ ,示意图如下：

    上图展示了如何将[144]的输入特征图经过一个33的滤波器滑动过程所对应的im2col转换过程，滤波器覆盖范围了的数其实im2col过程之后是不会连续存储在一起的，而是会以channel的形式进行折叠，每次滑动中的第一个元素会在col_buffer中连续存储。如果是多通道图像，由于连续存储的特性，它会保存在下一个存储块。
    在CUDA编程或者HIP编程中，开多线程来对这个转换过程进行并发，由于是多个线程进行同的并发过程，每个kernel函数中每个线程多对应的操作的是一样的，所以我们需要对每个线程要去做什么，多个线程之间能否一起协调的合作做出考虑。在设计过程中，每次起n个线程，每个线程具有自己的Index，这里可以看成一个线程池，每个线程在线程池中具有自己位置的索引。我们一起维护一个含有输出数据形状大小的数量的线程池，即 $[\frac{C_{out}}{g},H_{out},W_{out}]$ ，这个线程池中的线程输出和输出特征图的像素数量可以对应上，但是在滑动的过程中，操作数是大于这个数量的，从图6可以看出中间过程中的示意图，im2col实际转换的输出矩阵大小是 $[\frac{C_{in}}{g}*k*k,H_{out}*W_{out}]$ （看成四维的形式），也就是说一个线程需要负责k*k个像素的填写。

gemm

这里gemm操作是采用的cublas中或者rocblas的gemm计算函数，这里仅仅以FP32为例来对这个过程进行描述，在cuda中采用cublasSgemm()函数，早rocblas函数中采用rocblas_gemm_ex()函数。
阅读上面im2col过程可知一个输入特征图数据经过im2col的过程后得到一个 $[\frac{C_{in}}{g}*k*k,H_{out}*W_{out}]$ 大小的矩阵xdata，权重w $[\frac{C_{out}}{g},\frac{C_{in}}{g},k,k]$ 可以看成一个为 $[\frac{C_{out}}{g},\frac{C_{in}}{g}*k*k]$ 的二维矩阵，则w*xdata就可以得到输出矩阵 $[\frac{C_{out}}{g},H_{out}*W_{out}]$ ,在使用gemm函数的过程需要注意一下其中矩阵A和B的位置。
cuBlas中

cublasHandle_t handle = blas_handle();
cudaError_t status = cublasSgemm(handle, (TB ? CUBLAS_OP_T : CUBLAS_OP_N), 
            (TA ? CUBLAS_OP_T : CUBLAS_OP_N), N, M, K, &ALPHA, B_gpu, ldb, A_gpu, lda,                    &BETA, C_gpu, ldc);
            //A_gpu xdata
            //B_gpu  w

rocBlas中

ROCBLAS_RETURN_IF_ERROR(rocblasGemmHelper(
             RocblasHandle(),
             rocblas_operation_none,
             rocblas_operation_none,
             output_image_size,M/conv_attrs_.group, kernel_dim,
             &alpha,
             b, output_image_size,  //xdata
             a, kernel_dim,         //w
             &beta,
             c, output_image_size));
             
 // gemm
inline rocblas_status rocblasGemmHelper(rocblas_handle handle,
                                        rocblas_operation transa,
                                        rocblas_operation transb,
                                        int m, int n, int k,
                                        const float* alpha,
                                        const float* A, int lda,
                                        const float* B, int ldb,
                                        const float* beta,
                                        float* C, int ldc) {
  return rocblas_gemm_ex(handle,
                         transa,
                         transb,
                         m, n, k,
                         alpha,
                         A, rocblas_datatype_f32_r, lda,
                         B, rocblas_datatype_f32_r, ldb,
                         beta,
                         C, rocblas_datatype_f32_r, ldc,
                         C, rocblas_datatype_f32_r, ldc,
                         rocblas_datatype_f32_r,
                         rocblas_gemm_algo_standard, 0, 0);

代码设计

调用模块:

const size_t kernel_rank = kernel_shape.size();

const int64_t input_image_size = input_shape.Size();
const int64_t output_image_size = output_shape.Size();
const int64_t kernel_size = TensorShape(kernel_shape).Size();
const int64_t X_offset = C / conv_attrs_.group * input_image_size;
const int64_t Y_offset = Y->Shape().Size() / Y->Shape()[0] / conv_attrs_.group;
const int64_t W_offset = W->Shape().Size() / conv_attrs_.group;
const int64_t kernel_dim = C / conv_attrs_.group * kernel_size;
const int64_t single_col_buffer_size = kernel_dim * output_image_size;

  //展开IM2col过程所需要的临时变量
const int64_t col_buffer_size = kernel_dim  *output_image_size;
const int64_t im2col_X_offset = C  * input_image_size;

auto* col_data = alloc->Alloc(sizeof(HipT) * SafeInt<size_t>(col_buffer_size));
BufferUniquePtr col_buffer(col_data, BufferDeleter(std::move(alloc)));
auto* col_buffer_data =reinterpret_cast <HipT*>(col_buffer.get());//static_cast

const HipT zero = ToHipType<T>::FromFloat(0.f);
const float alpha = 1.0f;

for (int image_id = 0; image_id < N; ++image_id) {
  for (int group_id = 0; group_id < conv_attrs_.group; ++group_id) {
     auto *a = W->Data<float>() + (group_id) * W_offset; //W
     auto *b = col_buffer_data+(group_id)*single_col_buffer_size;
     auto *im = Xdata + (image_id*conv_attrs_.group+group_id)*X_offset; //X
     auto *c = Ydata + (image_id*conv_attrs_.group+group_id)*Y_offset;

     const float alpha = 1.0;
     const float beta = 0.0;

     ROCBLAS_RETURN_IF_ERROR(rocblasGemmHelper(
        RocblasHandle(),
        rocblas_operation_none,
        rocblas_operation_none,
        output_image_size,M/conv_attrs_.group, kernel_dim,
        &alpha,
        b, output_image_size,
        a, kernel_dim,
        &beta,
        c, output_image_size));
      }
    }

if (Bdata!=nullptr){
      //void add_bias_gpu(hipStream_t stream,T *output,const T *biases,const int batch,const int c_out,const int out_putsize)
      add_bias_gpu<HipT>(Stream(),Ydata, Bdata, static_cast<int>(N), static_cast<int>(M), static_cast<int>(output_image_size));
}

im2col模块

template <typename T>
__global__ void im2col_gpu_kernel(const int n, const T* data_im,
        const int height, const int width, const int ksize_h, const int ksize_w,
        const int pad_t, const int pad_l,const int pad_b,const int pad_r,
        const int stride_h,const int stride_w,const int dilation_h,const int dilation_w,
        const int output_h, const int output_w,
        T *data_col,const T padding_value
    ) {
    int index = blockIdx.x*blockDim.x+threadIdx.x;
    for(; index < n; index += blockDim.x*gridDim.x){//
        int w_out = index % output_w; //线程池的w_out
        int h_index = index / output_w;
        int h_out = h_index % output_h;
        int channel_in = h_index / output_h;
        int channel_out = channel_in * ksize_h * ksize_w;
        int h_in = h_out * stride_h - pad_t;
        int w_in = w_out * stride_w - pad_l;
        T* data_col_ptr = data_col;
        data_col_ptr += (channel_out * output_h + h_out) * output_w + w_out;
        const T* data_im_ptr = data_im;
        data_im_ptr += (channel_in * height + h_in) * width + w_in;
        for (int i = 0; i < ksize_h; ++i) {
            for (int j = 0; j < ksize_w; ++j) {
                int h = h_in + i*dilation_h;
                int w = w_in + j*dilation_w;

                *data_col_ptr = (h >= 0 && w >= 0 && h < height && w < width) ?
                    data_im_ptr[i*dilation_h * width + j*dilation_w] : padding_value;
                data_col_ptr += output_h * output_w;
            }
        }
    }
}

template <typename T>
void im2col_gpu(hipStream_t stream,const T *im,
         int channels, int height, int width,
         int ksize_h,int ksize_w, int stride_h,int stride_w, 
         int pad_t,int pad_l,int pad_b,int pad_r,
         int dilation_h,int dilation_w,
         T *data_col){
    // We are going to launch channels * height_col * width_col kernels, each
    // kernel responsible for copying a single-channel grid.
    const T padding_value = 0;
    int output_h = (height + pad_b + pad_t - (dilation_h * (ksize_h - 1) + 1)) / stride_h + 1;
    int output_w = (width + pad_l + pad_r - (dilation_w * (ksize_w - 1) + 1)) / stride_w + 1;
    int num_kernels = channels * output_h * output_w;
    im2col_gpu_kernel<<<(num_kernels+BLOCK-1)/BLOCK,
        BLOCK,0,stream>>>(
                num_kernels, im, height, width, ksize_h,ksize_w, pad_t,pad_l,pad_b,pad_r,
                stride_h,stride_w,dilation_h, dilation_w,
                output_h,output_w, data_col,padding_value);
    //if(hipDeviceSynchronize()) printf("hipDeviceSynchronize failed at im2col_gpu\n ");
    //printf("im2col kernel done\n ");
}

瓶颈与优化

上述代码其实有瓶颈，一是每次代码都需要做im2col的转存操作，增加了全局内存的读写时间；二是im2col过程中有两次for循环，一次是根据数据的batch_id,另外一个是根据分组数，在实际测速的时候发现在分组卷积的时候有瓶颈，回想这个设计的过程，im2col这个过程其实没有这个限制，主要是采用rocblas中gemm的时候不能多线程，一个是RocblasHandle()是单线程的，另一个是他们都会默认在0流上运行，如果需要多线程执行gemm就需要分流且多创建几个RocblasHandle()，然而RocblasHandle()最好是创建全局的，不至于是这个卷积计算的时候创建RocblasHandle()，然后这个卷积计算完后又释放掉，一个模型中含有很多卷积，如果一直这样创建RocblasHandle()然后释放RocblasHandle()，必然十分浪费时间。
解决这个问题，引入rocblas的批量处理，改进调用接口：

const size_t kernel_rank = kernel_shape.size();

const int64_t input_image_size = input_shape.Size();
const int64_t output_image_size = output_shape.Size();
const int64_t kernel_size = TensorShape(kernel_shape).Size();
const int64_t X_offset = C / conv_attrs_.group * input_image_size;
const int64_t Y_offset = Y->Shape().Size() / Y->Shape()[0] / conv_attrs_.group;
const int64_t W_offset = W->Shape().Size() / conv_attrs_.group;
const int64_t kernel_dim = C / conv_attrs_.group * kernel_size;
const int64_t single_col_buffer_size = kernel_dim * output_image_size;

//展开IM2col过程所需要的临时变量
const int64_t col_buffer_size = kernel_dim *conv_attrs_.group *output_image_size;
const int64_t im2col_X_offset = C  * input_image_size;

auto* col_data = alloc->Alloc(sizeof(HipT) * SafeInt<size_t>(col_buffer_size));
BufferUniquePtr col_buffer(col_data, BufferDeleter(std::move(alloc)));
auto* col_buffer_data =reinterpret_cast <HipT*>(col_buffer.get());//static_cast

const HipT zero = ToHipType<T>::FromFloat(0.f);
const float alpha = 1.0f;
//const float beta = 0.0f;

if(kernel_rank==2||kernel_rank==1){ 
for (int image_id = 0; image_id < N; ++image_id) {
  auto *temp_b = col_buffer_data;
  auto *im_src = reinterpret_cast<const HipT*>(Xdata + (image_id)*im2col_X_offset); //X

  if(kernel_rank==2)
     im2col_gpu<HipT>(Stream(),im_src, C ,input_shape[0],input_shape[1],kernel_shape[0],kernel_shape[1],strides[0],strides[1],pads[0],pads[1],pads[2],pads[3],dilations[0],dilations[1],temp_b);
  else if(kernel_rank==1)
     im2col_gpu<HipT>(Stream(),im_src, C ,1,input_shape[0],1,kernel_shape[0],1,strides[0],0,pads[0],0,pads[1],1,dilations[0],temp_b);//最后一个0是padding的value
     auto *a = Wdata ; //W
     auto *b = col_buffer_data;
     auto *c = Ydata + (image_id*conv_attrs_.group)*Y_offset;

     const int stride_A = M/conv_attrs_.group*kernel_dim;
     const int stride_B = output_image_size*kernel_dim;
     const int stride_C = M/conv_attrs_.group*output_image_size;

     ROCBLAS_RETURN_IF_ERROR(rocblasGemmStridedBatchedHelper(
            RocblasHandle(),
            rocblas_operation_none,
            rocblas_operation_none,
            static_cast<int>(output_image_size),static_cast<int>(M/conv_attrs_.group), static_cast<int>(kernel_dim),
            &alpha,
            b, static_cast<int>(output_image_size),stride_B,//x
            a, static_cast<int>(kernel_dim),stride_A, //w
            &alpha,
            c, static_cast<int>(output_image_size),stride_C,
            static_cast<int>(conv_attrs_.group)));
 }