技术心得:基于OneFlow实现Unfold、Fold算子

简介: 技术心得:基于OneFlow实现Unfold、Fold算子

1、从卷积层说起


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


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


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


from mxnet import autograd, nd


from mxnet.gluon import nn


def corr2d(X, K): # 本函数已保存在d2lzh包中方便以后使用


h, w = K.shape


Y = nd.zeros((X.shape【0】 - h + 1, X.shape【1】 - w + 1))


for i in range(Y.shape【0】):


for j in range(Y.shape【1】):


Y【i, j】 = (X【i: i + h, j: j + w】 K).sum()


return Y


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


2、初见img2col


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


图源来自微软AI-system仓库


这是微软的AISystem仓库对应的章节,强烈推荐大家去学习(本人鸽了好久没看完)


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


3、理解img2col


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


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


经过img2col后,输入特征图会变换为 (N, CinKK, OhOw) 这么一个三维向量。


此外卷积核我们会reshape成一个二维向量:(Cout, KKCin)。


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


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


4、img2col源码


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


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


def darknet_img2col(data, channels, height, width, ksize, stride, pad):


out_h = int((height + 2pad - ksize) / stride) + 1


out_w = int((width + 2pad - ksize) / stride) + 1


channels_cols = channelsksizeksize


out_shape = (channels_cols, out_hout_w)


elem_cnt = out_shape【0】 out_shape【1】


out_array = np.zeros(shape=elem_cnt, dtype=np.float32)


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


PyTorch的Unfold文档


而channel_cols则是之前我们提到的,img2col会把第二个维度变换为C_inKK。


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


for c in range(channels_cols):


# 分别计算一个kk的窗口内,h,w的偏移量


kh_offset = (c // ksize) % ksize


kw_offset = c % ksize


# 计算当前处理的通道index


c_im = c // ksize // ksize


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


for h in range(out_h):


for w in range(out_w):


im_row = kh_offset + h stride


im_col = kw_offset + w stride


index = (c out_h + h) out_w + w


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


取元素逻辑


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


当前channelOhOw + 当前hOw + 当前w


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


out_array【index】 = im2col_get_pixel(data, height, width, c_im, im_row, im_col, pad)


out_array = np.reshape(out_array, out_shape)


return out_array


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


def im2col_get_pixel(im, height, width, channel, row, col, pad):


row -= pad


col -= pad


if row < 0 or col < 0 or row

return 0


return im【int(col + width (row + height channel))】 # batchwhc + Width<span class="hljs-operator">Height<span class="hljs-operator">channel + Width<span class="hljs-operator">row + col


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


x = np.arange(1, 10).astype(np.float32)


out = darknet_img2col(x, channels=1, height=3, width=3, ksize=2, stride=1, pad=0)


输出结果符合预期:


【【1. 2. 4. 5.】


【2. 3. 5. 6.】


【4. 5. 7. 8.】


【5. 6. 8. 9.】】


col2img


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


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


OneFlow对应的实现


PR地址:


5、OneFlow版本的Unfold


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


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


输入内存排布


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


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


in_index_helper = NdIndexOffsetHelper[span class="hljs-type">INDEX_T, kInputNDim>(input_dims); // 格式为(N, C, H, W)


out_index_helper = NdIndexOffsetHelper[span class="hljs-type">INDEX_T, kOutputNDim>(output_dims); // 格式为(N, C, Kh, Kw, Oh, Ow)


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


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


// index_a format: (N, C, di, hi, wi, db, hb, wb) or (N, di, hi, wi, db, hb, wb, C)


// index_b format: (N, C, D, H, W) or (N, D, H, W, C)


// return: true indicates out-of-bound, otherwise in-bound


template[span class="hljs-keyword">typename INDEX_T, int NDIM, int SDIM>


OF_DEVICE_FUNC bool UnfoldIndexTransform(const UnfoldParams& params,


const INDEX_T index_a, INDEX_T index_b) {


// batch dim index transform


index_b【0】 = index_a【0】;


// channel dim index transform


using ParamType = UnfoldParams;


index_b【ParamType::kInputChannelDim】 = index_a【ParamType::kOutputChannelDim】;


// spatial dim index transform


// D,H,W spatial dim index transform


for (int64_t d = 0; d < NDIM; ++d) {


INDEX_T idx = index_a【SDIM + NDIM + d】 params.stride【d】


+ index_a【SDIM + d】 params.dilation【d】 - params.padding【d】;


if (idx < 0 || idx >= params.dims【d】) return true;


index_b【SDIM + d】 = idx;


}


return false;


}


模板参数 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为例):


Ohstride_h + khdilation_h - padding_h


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


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


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


template[span class="hljs-keyword">typename T, typename INDEX_T, int NDIM, int SDIM>


global void CudaUnfoldForward(UnfoldParams params, const T in, T out) {


CUDA_1D_KERNEL_LOOP_T(INDEX_T, out_offset, params.out_elem_cnt) {


//代码参考: https://weibo.com/u/7930938463

using ParamType = UnfoldParams;


INDEX_T in_index【ParamType::kInputNDim】 = {0};


INDEX_T out_index【ParamType::kOutputNDim】 = {0};


params.out_index_helper.OffsetToNdIndex(out_offset, out_index);


if (!UnfoldIndexTransform(params, out_index, in_index)) {


INDEX_T in_offset = params.in_index_helper.NdIndexToOffset(in_index);


out【out_offset】 = in【in_offset】;


} else {


out【out_offset】 = static_cast(kUnfoldPaddingValue);


}


}


}


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


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


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


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


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


6、OneFlow版本的Fold


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


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


template[span class="hljs-keyword">typename T, typename INDEX_T, int NDIM, int SDIM>


global void CudaFoldForward(FoldParams params, const T input_ptr,


T output_ptr) {


CUDA_1D_KERNEL_LOOP_T(INDEX_T, in_offset, params.in_elem_cnt) {


using ParamType = FoldParams;


INDEX_T in_index【ParamType::kInputNDim】 = {0};


INDEX_T out_index【ParamType::kOutputNDim】 = {0};


params.in_index_helper.OffsetToNdIndex(in_offset, in_index);


if (!FoldIndexTransform(params, in_index, out_index)) {


<div class="hljs-

相关文章
|
机器学习/深度学习 PyTorch Go
YOLOv5的Tricks | 【Trick4】参数重结构化(融合Conv+BatchNorm2d)
这篇文章是想要记录yolov5在模型搭建过程中的一个融合模块,就是把卷积与批归一化的参数进行融合,想卷积带有批归一化的性质,使得推理过程中可以加快模型推理速度,简化整个模型结构,实现训练与推理两个阶段的解耦。
930 0
YOLOv5的Tricks | 【Trick4】参数重结构化(融合Conv+BatchNorm2d)
|
6月前
|
数据可视化
R语言nlme、nlmer、lme4用(非)线性混合模型non-linear mixed model分析藻类数据实例(上)
R语言nlme、nlmer、lme4用(非)线性混合模型non-linear mixed model分析藻类数据实例
|
6月前
|
资源调度 安全
R语言nlme、nlmer、lme4用(非)线性混合模型non-linear mixed model分析藻类数据实例(下)
R语言nlme、nlmer、lme4用(非)线性混合模型non-linear mixed model分析藻类数据实例
|
6月前
|
机器学习/深度学习 分布式计算 数据库连接
[Spark精进]必须掌握的4个RDD算子之filter算子
[Spark精进]必须掌握的4个RDD算子之filter算子
144 2
|
并行计算 调度 Windows
R语言- parallel::mclapply 并行化计算任务
R语言中的 parallel::mclapply 是一个用于在多核CPU上实现并行计算的方法。它是lapply函数的并行版本,可以在多个处理器核心上同时运行lapply函数。mclapply函数的语法与lapply函数类似,但它可以指定要使用的处理器核心数量,从而提高计算速度。
1028 0
|
机器学习/深度学习 存储 并行计算
【Pytorch】Tensor的分块、变形、排序、极值与in-place操作
【Pytorch】Tensor的分块、变形、排序、极值与in-place操作
548 0
|
机器学习/深度学习 存储 人工智能
神经网络算子优化之——算子融合
算子融合,作为神经网络性能优化的一个必要手段,其性能收益是很高的。
|
机器学习/深度学习 人工智能 资源调度
深度学习入门基础CNN系列——批归一化(Batch Normalization)和丢弃法(dropout)
批归一化方法(Batch Normalization,BatchNorm)是由Ioffe和Szegedy于2015年提出的,已被广泛应用在深度学习中,其目的是对神经网络中间层的输出进行标准化处理,使得中间层的输出更加稳定。丢弃法(Dropout)是深度学习中一种常用的抑制过拟合的方法,其做法是在神经网络学习过程中,随机删除一部分神经元。训练时,随机选出一部分神经元,将其输出设置为0,这些神经元将不对外传递信号。
546 0
深度学习入门基础CNN系列——批归一化(Batch Normalization)和丢弃法(dropout)
|
分布式计算 算法 大数据
Rdd 算子_转换_sample | 学习笔记
快速学习 Rdd 算子_转换_sample
154 0
Rdd 算子_转换_sample | 学习笔记
|
机器学习/深度学习 计算机视觉
Batch Normalization, 批标准化,神经网络shortcut 是什么,无脑用ReLU(CV领域). 无脑用3x3.
和普通的数据标准化类似, 是将分散的数据统一的一种做法, 也是优化神经网络的一种方法. 在之前 Normalization 的简介视频中我们一提到, 具有统一规格的数据, 能让机器学习更容易学习到数据之中的规律.
380 0