Skip to main content

moregeek program

实录 | megengine 大 kernel 卷积工程优化实践-多极客编程


本篇整理自 3.19 日 “Large Kernel Makes CNN Great Again” 专题 Meetup 中《MegEngine 大 Kernel 卷积工程优化实践》分享。


视频回顾 01:10:55 直达



从卷积到矩阵乘


矩阵乘(GEMM)具有计算密度大和易并行等优良特点,传统领域例如芯片产业和高性能计算领域往往将 GEMM 作为 benchmark 并已经将其优化到接近硬件理论峰值。 为了获得更好的性能收益,im2col 算法将 GEMM 带进了卷积神经网络的工程优化领域。Implicit GEMM 算法进一步解决了 im2col 固有的多余显存占用和冗余的前后处理问题,这在存储受限的硬件例如 GPU 上尤为重要,使得 GEMM 在卷积优化中的重要性进一步提升。加之硬件厂商也开始越来越多的对矩阵乘提供硬件支持,例如各种 MMA 指令和 TensorCore 的加入,这些原因共同促使现有很多优化算法库已经将 im2col/Implicit GEMM 作为其默认的卷积优化方案。


im2col 算法


本文中假设卷积的输入 shape 为 (n, ic, ih, iw),kernel 为 (oc, ic, kh, kw),output 为 (n, oc, oh, ow)。im2col 算法的过程如下图所示,简单的将 kernel reshape 一下就得到了一个行 M = oc,列 K = ic*kh*kw 的矩阵记作矩阵 A。用 kernel 大小的立方体在 input 上做滑窗,每次将一个小立方体的数据按照 chw 的顺序展开成一列。从上到下,从左到右滑完整个 input 之后将会得到一个行 K = ic*kh*kw,列 N = n*oh*ow 的矩阵记作矩阵 B。此时我们计算 GEMM(A, B) 就可以得到卷积的结果矩阵 C,其行 M = oc,列 N = n*oh*ow。



在 GEMM 的计算过程中,根据 m, n 和 k 三个维度下标可以推导出数据在卷积的输入输出以及 kernel 中的下标。所以可以将 im2col 和 GEMM 两个过程融合在一起从而达到降低显存占用和性能加速的效果,这其实就是 Implicit GEMM 的原理。本文不过多介绍,感兴趣的可以阅读之前的技术文章


Implicit Batched GEMM


上一篇文章主要介绍了 MegEngine 大 kernel depthwise 卷积优化的背景和动机,本篇文章将介绍具体的优化思路和工程实践。借助 im2col/Implicit GEMM 算法,GEMM 在传统的针对 dense 卷积的优化中表现出来了优良的性能。所以针对大 kernel depthwise 卷积也应该尝试使用 GEMM 实现。如前文分析,直接使用和 dense 卷积一样的方法将 im2col 算法应用到大 kernel depthwise 卷积将会产生一个 Batched GEMV,这很难达到硬件浮点计算峰值。显然为了将 GEMM 应用在大 kernel depthwise 卷积上,我们应该转变一下思维。如下图所示,dense 卷积一般用 kernel 去卷 input 并算得 output,但当 kernel size 很大甚至可能比 input 还大时,此时其实应该用 input 去卷 kernel 并算得 output。



回想一下 dense 卷积用 kernel 去卷 input,对 input 做 im2col 变换。现在大 kernel depthwise 卷积用 input 去卷 kernel,所以此时应该对 kernel 做 im2col 变换。算法过程没有本质区别,只需要在 im2col 中将 kernel 看做 input,将 input 看做 kernel 即可。由于 depthwise 卷积是逐通道做卷积的,因此 im2col 变换也需要逐通道做。如下图所示,每个通道变换之后都会产生一个 M = n, N = oh*ow, K = ih*iw 的 GEMM。根据上一篇文章的分析,Batched GEMM 相比于 Batched GEMV 更容易打满硬件设备的浮点计算峰值。



CUTLASS 是 NVIDIA 的开源模板库,它旨在提供一种用较小的成本写出一个性能不是那么差的 GEMM 的能力。CUTLASS 内置了针对 GEMM 的 meta schedule,能够让计算尽量掩盖访存延迟从而达到不错的性能。旷视早在 CUTLASS 官方开源其卷积实现之前就基于 CUTLASS 做出了自己的卷积实现,时至今日已经打磨出了一个更适合内部业务的旷视版 CUTLASS。此处的 Implicit Batched GEMM 也是基于旷视版 CUTLASS 实现的,代码已经随 MegEngine v1.8.2 开源出来了,实现细节就不过多介绍了。如下图的实验数据显示随着 kernel size 的增加,Implicit Batched GEMM 的性能大致是呈线性增长的,部分情况下可以逼近理论峰值。



Implicit Batched GEMM 的优点一方面是可以复用成熟的 GEMM 优化思路和基础设施,还可以方便使用 TensorCore 进行加速;另一方面如果在推理的时候不要求可变 shape 的话,对 kernel 的 im2col 变换可以提前算好进一步加速。当然它的缺点也很明显,比如小 batch 情况下依然会退化成 Bacthed GEMV。如果用 M*N*K*2 来近似 GEMM 的计算量的话,不难发现 Implicit Batched GEMM 的计算量相比 dense 卷积转成的 GEMM 增大了 $\frac{ihiw}{khkw}$ 倍,这意味着 Implicit Batched GEMM 在 input 显著大于 kernel size 时性能不佳。如下图所示的实验结果也显示着当 input 大于 kernel size 时,随着 input 的增加 Implicit Batched GEMM 的性能有明显下滑。需要一种新的优化方法来迎合下游如检测、分割等业务里的大 input size 的需求,这种方法在小 batch size 或者大输入下的性能表现也要足够好。



Direct Conv


由于大 kernel depthwise 卷积的计算密度比较高,所以其实简单实现一版性能基本都能达到峰值性能的 70%-80%。Driect Conv 的写法其实有很多,这里只提供一种写法思路供参考。如下图所示,为了更好的利用 CUDA 的多级存储以最大利用带宽,Direct Conv 采用多级分块策略。每个 Thread Block 负责计算 output 的一个分块,然后每个 Warp 对 Thread Block Tile 按行进一步做分块。为了适应更大的 kernel size,我们在 Thread level 上不仅针对 output 做了分块,还对 kernel 做了分块。



简单举个例子介绍 Thread level 的分块策略。假设 Thread Block size 是 128,Thread 被组织成 32×4 的形式,每一行的 4 个线程负责计算 output 的一行。将 kernel 也切分成四列,每一行的 4 个线程分别负责读取 kernel 的一列。如下图所示,Thread 0 读取 kernel 的第 0 列和 input 的第 0-3 列,计算得到 4 个 output;Thread 1 读取 kernel 的第 2 列和 input 的第 1-4 列,计算得到 4 个 output。Thread 2 和 Thread 3 以此类推。



由于对 kernel 做了分块,所以每行的 4 个线程计算完毕之后每个 Thread 持有的是 output 的部分和,需要将 4 个线程各自的结果规约到一起才是最终结果。此处借助了 Warp Shuffle API__shuffle_xor_sync,它实现了一种蝶形规约,其原理如下图所示。由于只需要将每 4 个线程的结果规约到一起就行了,所以只需要进行 2 次 __shuffle_xor_sync 即可,最后将 outupt 写回。



实验数据显示在 input 大小为 48 时 Direct Conv 的性能已经略高于 Implicit Batched GEMM 了,intput 为 64 时 Direct Conv 的性能会显著高于 Implicit Batched GEMM。得益于 MegEngine 的算子自动选择机制,用户使用的时候不用指定具体的实现方式,MegEngine 会自动选择最佳实现。



运行时间


为了衡量算子的优劣,前面的实验都是从算子绝对性能和硬件理论峰值相比的角度设计的。为了让用户有更直观的感受,我们同样测试了大 kernel depthwise 卷积的运行时间。实验环境为 2080Ti @ cuda10.1 + cudnn7.6.3,所用的数据类型为 fp32, batch size 为 64,channel 为 384,用 24 个 layer 进行前向和反向计算。从下图可见 MegEngine 比 PyTorch(with cudnn) 最高快 10 倍以上,优化后的 MegEngine 在 31×31 的 kernel size 上和 PyTorch 9×9 的训练时间相当。



如下图所示,只测试一个 layer 的前向推理,其他的配置和训练保持一致。经过优化后的 MegEngine 比 cudnn 最高快 8 倍多,并且 fp16 相比 fp32 也有 2 倍多加速,欢迎尝试一下混合精度训练。代码已经随着 MegEngine v1.8.2 开源,使用 v1.9 (即将发布)效果更佳~




更多 MegEngine 信息获取,您可以:查看文档、和 GitHub 项目


©著作权归作者所有:来自51CTO博客作者MegEngineBot的原创作品,请联系作者获取转载授权,否则将追究法律责任

机器学习技术系列:一篇图文笔记了解【机器学习基础知识】-多极客编程

导言 最近有小半年由近半数工作和生活时间在机器学习技术(ML)的学习与工程实践中,感觉自己阅读了几本ML方面好书,找到了一些更好的学习网站,所以重新梳理了一下自己理解的的ML基础知识。 相关参考摘录书籍及网站如下 《机器学习实战:基于Scikit-Learn、Keras和TensorFlow》(第2版) 《Python深度学习》(第2版) 网站:https://www.showmeai.tech

深度学习背景下的图像三维重建技术进展综述-多极客编程

原文首发于《中国图象图形学报》作者:杨航,陈瑞,安仕鹏,魏豪,张衡原文地址:​​深度学习背景下的图像三维重建技术进展综述​​三维重建是指从单张二维图像或多张二维图像中重建出物体的三维模型,并对三维模型进行纹理映射的过程。三维重建可获取从任意视角观测并具有色彩纹理的三维模型,是计算机视觉领域的一个重要研究方向。传统的三维重建方法通常需要输入大量图像,并进行相机参数估计、密集点云重建、表面重建和纹理映

模型压缩部署概述-多极客编程

一,模型在线部署 深度学习和计算机视觉方向除了算法训练/研究,还有两个重要的方向: 模型压缩(模型优化、量化)、模型部署(模型转换、后端功能SDK开发)。所谓模型部署,即将算法研究员训练出的模型部署到具体的端边云芯片平台上,并完成特定业务的视频结构化应用开发。 现阶段的平台主要分为云平台(如英伟达 GPU)、手机移动端平台(ARM 系列芯片)和其他嵌入式端侧平台(海思 3519、安霸 CV22、地

目标检测模型的评价标准-ap与map-多极客编程

前言 一,精确率、召回率与F1 1.1,准确率 1.2,精确率、召回率 1.3,F1 分数 1.4,PR 曲线 1.4.1,如何理解 P-R 曲线 1.5,ROC 曲线与 AUC 面积 二,AP 与 mAP 2.1,AP 与 mAP 指标理解 2.2,近似计算AP 2.3,插值计算 AP 2.4,mAP 计算方法 三,目标检测度量标准汇总 四,参考资料 前言 为了了解模型的泛

onnx模型分析与使用-多极客编程

本文大部分内容为对 ONNX 官方资料的总结和翻译,部分知识点参考网上质量高的博客。 一,ONNX 概述 深度学习算法大多通过计算数据流图来完成神经网络的深度学习过程。 一些框架(例如CNTK,Caffe2,Theano和TensorFlow)使用静态图形,而其他框架(例如 PyTorch 和 Chainer)使用动态图形。 但是这些框架都提供了接口,使开发人员可以轻松构建计算图和运行时,以优

基于单机最高能效270亿参数gpt模型的文本生成与理解-多极客编程

作者:李鹏,王玮,陈嘉乐,黄松芳,黄俊单位:阿里云智能机器学习平台PAI & 达摩院自然语言基础技术概述GPT模型能较好的处理文本生成领域的各种任务,比如文本补全,自由问答,完形填空,写作文,写摘要,写小说,写诗歌等等。最近火爆全网的人工智能产品ChatGPT也是以GPT文本生成模型为底座。虽然GPT大模型作用在这些应用领域的效果很好,但是训练成本非常高。以OpenAI推出的1750亿的G

模型压缩部署概述-多极客编程

一,模型在线部署 深度学习和计算机视觉方向除了算法训练/研究,还有两个重要的方向: 模型压缩(模型优化、量化)、模型部署(模型转换、后端功能SDK开发)。所谓模型部署,即将算法研究员训练出的模型部署到具体的端边云芯片平台上,并完成特定业务的视频结构化应用开发。 现阶段的平台主要分为云平台(如英伟达 GPU)、手机移动端平台(ARM 系列芯片)和其他嵌入式端侧平台(海思 3519、安霸 CV22、地

arm 算子性能优化上手指南-多极客编程

作者:林锦豪 | 旷视 MegEngine 架构师 前言 做 arm 侧算子开发时,不能不关心的就是性能。本文主要就是介绍 arm 算子性能优化的常用思路,做为一个入门级的参考。文章以 ARM Cortex a55 上的 GaussianBlur 优化为例展开,并在文末对 arm 性能优化思路做了一个总结。 GaussianBlur 的优化 Q1: 什么是 GaussianBlur? Gaus

github 车牌检测识别项目调研-多极客编程

一,EasyOCR 1.1,仓库介绍 1.2,使用记录 二,HyperLPR 2.1,HyperLPR 概述 2.3,使用记录 2.3,使用建议 三,simple-car-plate-recognition-2 3.1,仓库介绍 3.2,使用记录 3.3,使用建议 四,车牌检测-License-Plate-Detector 4.1,仓库介绍 4.2,建议 五,MMOCR

目标检测模型的评价标准-ap与map-多极客编程

前言 一,精确率、召回率与F1 1.1,准确率 1.2,精确率、召回率 1.3,F1 分数 1.4,PR 曲线 1.4.1,如何理解 P-R 曲线 1.5,ROC 曲线与 AUC 面积 二,AP 与 mAP 2.1,AP 与 mAP 指标理解 2.2,近似计算AP 2.3,插值计算 AP 2.4,mAP 计算方法 三,目标检测度量标准汇总 四,参考资料 前言 为了了解模型的泛

onnx模型分析与使用-多极客编程

本文大部分内容为对 ONNX 官方资料的总结和翻译,部分知识点参考网上质量高的博客。 一,ONNX 概述 深度学习算法大多通过计算数据流图来完成神经网络的深度学习过程。 一些框架(例如CNTK,Caffe2,Theano和TensorFlow)使用静态图形,而其他框架(例如 PyTorch 和 Chainer)使用动态图形。 但是这些框架都提供了接口,使开发人员可以轻松构建计算图和运行时,以优

基于单机最高能效270亿参数gpt模型的文本生成与理解-多极客编程

作者:李鹏,王玮,陈嘉乐,黄松芳,黄俊单位:阿里云智能机器学习平台PAI & 达摩院自然语言基础技术概述GPT模型能较好的处理文本生成领域的各种任务,比如文本补全,自由问答,完形填空,写作文,写摘要,写小说,写诗歌等等。最近火爆全网的人工智能产品ChatGPT也是以GPT文本生成模型为底座。虽然GPT大模型作用在这些应用领域的效果很好,但是训练成本非常高。以OpenAI推出的1750亿的G