• 基于OneFlow实现Unfold、Fold算子

    0825036e622c8fa00f0f58054ea77fa2.png

    撰文 | zzk

    1

    从卷积层说起

    熟悉CNN的小伙伴应该知道卷积是一个很常用也很重要的操作，CNN里的卷积和信号处理的卷积并不是一回事，CNN的卷积是做一种二维的互相关运算，以《动手学深度学习》5.1章为示例：

    922f043933408de0feafbe4197c28897.png

    《动手学深度学习》5.1.1. 二维互相关运算

    窗口内的元素和卷积核相乘，求和得到输出元素，一份naive的代码如下（同来自《动手学深度学习》

    1. from mxnet import autograd, nd
    2. from mxnet.gluon import nn
    3.  
    4. def corr2d(X, K):  # 本函数已保存在d2lzh包中方便以后使用
    5.     h, w = K.shape
    6.     Y = nd.zeros((X.shape[0] - h + 1, X.shape[1] - w + 1))
    7.     for i in range(Y.shape[0]):
    8.         for j in range(Y.shape[1]):
    9.             Y[i, j] = (X[i: i + h, j: j + w] * K).sum()
    10.     return Y

    这里是借助了numpy array的索引特性来写的，如果在c++里写，需要的循环层数更多（会在后面进行展示）。而这种循环计算的写法效率并不高，会极大拖慢卷积运算的速度。

    2

    初见img2col

    为了提高卷积运算的速度，img2col算法被发明了出来的，它的本质是用矩阵乘法来等价卷积运算，示例图如下：

    6936416baaa2392d478bfe47dd971b14.png

    图源来自微软AI-system仓库

    https://github.com/microsoft/AI-System/blob/main/docs/SystemforAI-4-Computer%20architecture%20for%20Matrix%20computation.pdf 这是微软的AISystem仓库对应的章节，强烈推荐大家去学习（本人鸽了好久没看完）

    可以看到img2col把输入特征图进一步展开，然后把filter展平，两者做矩阵乘，得到了相同的结果。

    3

    理解img2col

    看了上面的图可能还是不能理解这是怎么展开的，这里我们会介绍下：

    • 假设输入特征图为(N, Cin, H, W)，卷积核参数为(K, K, Cin, Cout), 输出特征图的长宽为(Oh, Ow)

    经过img2col后，输入特征图会变换为 (N, Cin*K*K, Oh*Ow) 这么一个三维向量。

    此外卷积核我们会reshape成一个二维向量：(Cout, K*K*Cin)。

    而这两个向量可以做一个矩阵乘，输出向量为(N, Cout, Oh*Ow)，这也是我们预期的卷积结果。

    img2col算法是一种用空间换时间的做法，将输入进行变换后，显然它所占用的显存空间变大了，好处则是可以借助矩阵乘，快速地完成卷积运算。

    下面我会结合darknet的原生img2col和一篇博客来进行相关讲解

    代码：

    https://github.com/pjreddie/darknet/blob/master/src/im2col.c

    博客：

    https://blog.csdn.net/caicaiatnbu/article/details/100515321

    4

    img2col源码

    darknet的img2col其实是搬运自caffe的，我们这里为了方便理解，以简单的CPU版本为例子，且不考虑batch维度。

    为了让读者能够快速运行上代码，这里讲解我以自己写的一版darknet img2col来作为示例。

    1. def darknet_img2col(data, channels, height, width, ksize, stride, pad):
    2.     out_h = int((height + 2*pad - ksize) / stride) + 1
    3.     out_w = int((width + 2*pad - ksize) / stride) + 1
    4.  
    5.     channels_cols = channels*ksize*ksize
    6.     out_shape = (channels_cols, out_h*out_w)
    7.     elem_cnt = out_shape[0] * out_shape[1]
    8.  
    9.     out_array = np.zeros(shape=elem_cnt, dtype=np.float32)

    首先我们可以确定输出tensor的各个维度，其中out_h和out_w就是输出的高，宽，采用的是卷积输出的公式：

    ea75fb7c22da9e95590be0c22ea30530.png

    PyTorch的Unfold文档

    而channel_cols则是之前我们提到的，img2col会把第二个维度变换为C_in*K*K。

    然后进入到次数为channel_cols的for循环

    1. for c in range(channels_cols):
    2.         # 分别计算一个k*k的窗口内，h，w的偏移量
    3.         kh_offset = (c // ksize) % ksize
    4.         kw_offset = c % ksize
    5.         # 计算当前处理的通道index
    6.         c_im = c // ksize // ksize

    然后我们需要根据当前处理的输出元素index，来获取对应输入的元素

    1. for h in range(out_h):
    2.             for w in range(out_w):
    3.                 im_row = kh_offset + h * stride
    4.                 im_col = kw_offset + w * stride
    5.                 index = (c * out_h + h) * out_w + w

    im_row的计算方式逻辑是：当前处理的输入元素窗口起始点：即h*stride，然后加上窗口内的kh_offset偏移量。

    1d11478d57138bca700c2fc9389eca99.png

    取元素逻辑

    而index的计算比较容易，因为输出是(C, H, W)，对应的一维index那就是

    当前channel*Oh*Ow + 当前h*Ow + 当前w

    最后我们再将元素取出来，赋给out数组。然后再将一维的out数组reshape成我们前面推导得到的out_shape

    1. out_array[index] = im2col_get_pixel(data, height, width, c_im, im_row, im_col,  pad)
    2.  
    3.     out_array = np.reshape(out_array, out_shape)
    4.  
    5.     return out_array

    img2col_get_pixel是一个合法取元素的函数，如果出现越界范围（比如小于0，或者大于Oh），那么就是取到padding的部分，此时我们返回0。

    1. def im2col_get_pixel(im, height, width, channel, row, col, pad):
    2.     row -= pad
    3.     col -= pad
    4.     if row < 0 or col < 0 or row >= height or col >= width:
    5.         return 0
    6.     return im[int(col + width * (row + height * channel))] # batch*w*h*c + width*height*channel + width*row + col

    我们可以简单构造一个数组来验证结果（以微软AI-System课程的示例作为输入）

    1. x = np.arange(1, 10).astype(np.float32)
    2. out = darknet_img2col(x, channels=1, height=3, width=3, ksize=2, stride=1, pad=0)

    输出结果符合预期：

    1. [[1. 2. 4. 5.]
    2.  [2. 3. 5. 6.]
    3.  [4. 5. 7. 8.]
    4.  [5. 6. 8. 9.]]

    col2img

    col2img则是img2col的逆过程，有兴趣的读者可以参考下这个博客：

    https://blog.csdn.net/caicaiatnbu/article/details/102626135

    在后面的oneflow实现里会有更完整的图例用于理解

    OneFlow对应的实现

    PR地址：

    https://github.com/Oneflow-Inc/oneflow/pull/5675

    5

    OneFlow版本的Unfold

    在深度学习框架里，img2col和col2img在Pytorch里还有另外的名字，也就是Unfold和Fold。通常想自己自定义一些跟卷积相关的骚操作时候，就经常会用到这两个Op。

    我们假设输入是一个(1, 2, 4, 4)的张量，但在框架内部，我们通常都是以一个一维数组来存储的，如下图所示：

    9a101bae6e8a64ab6158017d7f6eb38e.png

    输入内存排布

    然而我们需要对应的高维数组索引，OneFlow内部有一个NdIndexHelper类，构造的时候我们可以传入高维shape，然后调用OffsetToNdIndex来进行一维offset到高维索引的转换。

    这里我们分别给输入，输出Tensor构造一个NdIndexHelper

    1. in_index_helper = NdIndexOffsetHelper(input_dims); // 格式为(N, C, H, W)
    2. out_index_helper = NdIndexOffsetHelper(output_dims); // 格式为(N, C, Kh, Kw, Oh, Ow)

    比较特别的是，我们把输出构造成一个6维的形式。

    接着就是从输出的视角来推导应该取哪一个输入元素

    1. // index_a format: (N, C, di, hi, wi, db, hb, wb) or (N, di, hi, wi, db, hb, wb, C)
    2. // index_b format: (N, C, D, H, W) or (N, D, H, W, C)
    3. // return: true indicates out-of-bound, otherwise in-bound
    4. templateint NDIM, int SDIM>
    5. OF_DEVICE_FUNC bool UnfoldIndexTransform(const UnfoldParams& params,
    6.                                          const INDEX_T* index_a, INDEX_T* index_b) {
    7.   // batch dim index transform
    8.   index_b[0] = index_a[0];
    9.   // channel dim index transform
    10.   using ParamType = UnfoldParams;
    11.   index_b[ParamType::kInputChannelDim] = index_a[ParamType::kOutputChannelDim];
    12. // spatial dim index transform
    13.  
    14.   // D,H,W spatial dim index transform
    15.   for (int64_t d = 0; d < NDIM; ++d) {
    16.     INDEX_T idx = index_a[SDIM + NDIM + d] * params.stride[d]
    17.                   + index_a[SDIM + d] * params.dilation[d] - params.padding[d];
    18.     if (idx < 0 || idx >= params.dims[d]) return true;
    19.     index_b[SDIM + d] = idx;
    20.   }
    21.   return false;
    22. }

    • 模板参数 INDEX_T表示Index的数据类型（可以有int32_t, int64_t），NDIM表示处理几维数据（这里我们是2维），SDIM则是决定通道维所在位置，SDIM=1是NHWC格式，SDIM=2则是NCHW格式（这里我们取2）

    • 输入参数 index_a表示输出的NdIndexHelper，index_b则表示的是输入的NdIndexHelper

    从前面我们可以看到N，C这两个维度的index是不变的，所以我们直接给过去

    然后进入一个次数为NDIM==2的循环

    这里index的计算是从输出往输入推，公式是（以H为例）：

    Oh*stride_h + kh*dilation_h - padding_h

    计算得到输入的index，如果小于0或者大于输入的宽高，则说明是到了padding的地方，我们直接return true，以表示越界。如果能取到元素，我们则将这个index赋给index_b即输入的NdIndexHelper，且return false。

    这部分的分解操作可参考下图：

    01f66aea111bf6da40ad65847c322bdd.png

    从输出推导的好处就是整个运算是一个elementwise的操作，我们可以用输出tensor的元素个数做一个循环完成整个unfold操作。

    1. templateint NDIM, int SDIM>
    2. __global__ void CudaUnfoldForward(UnfoldParams params, const T* in, T* out) {
    3.   CUDA_1D_KERNEL_LOOP_T(INDEX_T, out_offset, params.out_elem_cnt) {
    4.     using ParamType = UnfoldParams;
    5.     INDEX_T in_index[ParamType::kInputNDim] = {0};
    6.     INDEX_T out_index[ParamType::kOutputNDim] = {0};
    7.     params.out_index_helper.OffsetToNdIndex(out_offset, out_index);
    8.     if (!UnfoldIndexTransform(params, out_index, in_index)) {
    9.       INDEX_T in_offset = params.in_index_helper.NdIndexToOffset(in_index);
    10.       out[out_offset] = in[in_offset];
    11.     } else {
    12.       out[out_offset] = static_cast(kUnfoldPaddingValue);
    13.     }
    14.   }
    15. }

    • 首先根据offset来计算当前处理输出元素的NdIndex

    • 然后判断UnfoldIndexTransform该方法的返回值

    • 如果为false，则说明我们可以取到输入元素，将其index转换为1d的offset，并赋值给输出

    • 如果为true，则越界，我们填充先前设定好的一个padding_value(0)

    至此整个img2col流程已经完成，整体操作可参考下图：

    aadd05ce76125e9577699c1e49efe8ef.png

    6

    OneFlow版本的Fold

    dc86cbc8f098ccc43a47137b69878803.png

    Fold就是将每一列填充回到kxk的地方

    如果能理解前面的Unfold，那么这里的Fold也能很容易的理解。只是操作的元素反过来而已。

    1. templateint NDIM, int SDIM>
    2. __global__ void CudaFoldForward(FoldParams params, const T* input_ptr,
    3.                                 T* output_ptr) {
    4.   CUDA_1D_KERNEL_LOOP_T(INDEX_T, in_offset, params.in_elem_cnt) {
    5.     using ParamType = FoldParams;
    6.     INDEX_T in_index[ParamType::kInputNDim] = {0};
    7.     INDEX_T out_index[ParamType::kOutputNDim] = {0};
    8.     params.in_index_helper.OffsetToNdIndex(in_offset, in_index);
    9.     if (!FoldIndexTransform(params, in_index, out_index)) {
    10.       INDEX_T out_offset = params.out_index_helper.NdIndexToOffset(out_index);
    11.       XPUAdd::Invoke(&input_ptr[in_offset], &output_ptr[out_offset]);
    12.     } else {
    13.       continue;
    14.     }
    15.   }
    16. }

    沿用前面的索引映射逻辑，我们进入一个循环次数为输入元素个数的循环体，计算得到当前offset对应的输入NdIndex。

    如果FoldIndexTransform返回的是false，则计算输出的offset，并使用原子加atomic add，把输入元素累加到该输出位置。

    7

    小结

    这部分代码是接手同事写一半的代码完成的，不得不说同事的构思真的很巧妙，通过模板参数能够拓展1d 2d 3d，nchw, nhwc各种格式，尽管直观上不太好理解。

    darknet版本(Caffe同款)是比较直观的帮助新手入门的img2col算法，可以结合那两篇CSDN博客来理解整个过程。

    其他人都在看

    点击“阅读原文”，欢迎下载体验OneFlow新一代开源深度学习框架

    0d88616eb47dd59b6d56353a7e3caeca.png

  • 相关阅读:
    【Java并发】聊聊如何优雅的暂停一个线程以及LockSupport
    技术人必读的「CDP选型指南」来了！
    Chapter 3. Char Drivers
    【Android 屏幕适配】屏幕适配基础概念 ① ( Android 与 iOS 屏幕宽高比种类 | 屏幕像素密度 DPI )
    SpringBoot入门
    从零开始，开发一个 Web Office 套件（16）：拖动控制点，调整编辑器大小
    Qt信号槽与事件循环学习笔记
    【蓝桥】小蓝的疑问
    postgresql源码学习（36）—— 事务日志11 - 日志归档
    python requests爬虫 POST无参请求一直在请求中，没有响应的原因
  • 原文地址：https://blog.csdn.net/OneFlow_Official/article/details/120735710