深度学习,这一现代图像识别和自动驾驶背后的技术,高度依赖于一种被称为“卷积”的数学运算。你可以将这种运算想象成一个在图像上滑动的窗口,它通过将小的像素块与一组模式进行比对,来识别边缘或纹理等特征。这个过程是人工智能的核心引擎,消耗了运行这些系统所需的大部分时间和能量。为了使这些系统更快、更高效,研究人员多年来一直致力于优化这种滑动窗口在计算机芯片上的移动方式,特别是在高端计算机中常见的强大图形处理器(GPU)上。挑战始终在于一种权衡:快速的方法往往需要大量的临时内存,而节省内存的方法则往往速度较慢或不稳定。
一支研究团队现在开发出了一种执行这些计算的新方法,打破了这种权衡。他们改进了一种称为“im2win”的技术,该技术重新组织了计算机在扫描过程中存储和访问图像数据的方式。通过改变数据的布局,新方法允许计算机以平滑、连续的流式方式移动信息,而不是以零散跳跃的方式。这种简单的组织形式的变化,使得计算机在处理过程中所需的临时内存比旧有的标准方法减少了一半以上。此外,研究人员还调整了这项技术,使其能够适配现代芯片中的两种不同类型的处理能力:处理精确计算的通用核心,以及旨在加速大规模数值块处理的专用“张量核心”(tensor cores)。
研究人员在十二种不同类型的图像处理任务上测试了他们的方法,涵盖了从简单滤波器到高级神经网络中复杂层的各种任务。他们发现,这种优化的方法明显快于目前的行业标准。当在专用的张量核心上运行时,新方法的性能达到了现有最佳软件库的1.4倍,以及一种基于矩阵乘法的常见替代方案的6.4倍。在速度方面(以每秒万亿次运算衡量),该方法比其在通用核心上运行的版本快了近三倍。或许最重要的一点是,这种速度提升是以大幅降低内存占用为代价的。新方法仅需常用矩阵法所需内存的35%,也仅为领先行业软件所用内存的53%。
为了实现这些结果,该团队引入了几项具体的工程改进。他们设计了以“之字形”(zig-zag)模式访问数据,这可以防止计算机内存的不同部分发生拥挤并相互拖慢速度。他们还建立了一个系统,让计算机在计算当前批次数据的同时,预先将下一批数据移动到位,从而有效地隐藏了移动信息所花费的时间。通过仔细测试,他们发现这种“双缓冲”(double buffering)技术是其速度提升中最关键的因素。虽然“之字形”模式有所帮助,但其重要性不及将数据移动与计算重叠的能力。
这项研究证实,通过将软件与计算机内存及处理单元的物理布局进行精细对齐,可以在不牺牲准确性的情况下,使深度学习系统变得更快且更具内存效率。研究人员证明,他们的方法在各种图像尺寸和滤波器形状下都能稳定运行,证明了它是应对现代人工智能多样化需求的稳健解决方案。通过解决内存开销和数据访问效率低下的问题,这项工作提供了一个统一的框架,有望让未来的AI系统能够在现有硬件上运行更复杂的模型,或者让当前的模型以显著更少的能量和时间运行。
技术摘要:在 GPU CUDA 和 Tensor Cores 上实现支持多精度的内存高效型 Im2win 卷积
问题陈述
卷积操作构成了深度神经网络(DNN)中的主要计算瓶颈,通常占据执行时间的 50%–90%。现有的 GPU 卷积方法面临显著的权衡:
- 直接卷积 (Direct Convolution): 提供零内存开销,但缓存利用率较低。
- 基于 FFT 的卷积 (FFT-based Convolution): 为大尺寸卷积核提供加速,但会引入小尺寸滤波器的延迟。
- Winograd 卷积: 减少了固定小尺寸卷积核的计算量,但在较大尺寸下表现出数值不稳定性。
- 基于 Im2col 的 GEMM: 利用了如 cuBLAS 等高度优化的库,但由于数据变换会产生高额的内存开销,并创建出导致性能欠佳的非规则矩阵。
- 隐式 GEMM (cuDNN): 虽然高效,但通常依赖于小批量(mini-batch)策略,这可能无法在所有卷积核尺寸和精度水平下充分挖掘硬件潜力。
核心挑战在于开发一种统一的卷积框架,以最小化内存开销、最大化缓存利用率、支持各种卷积核尺寸,并高效利用全精度 CUDA 核心和半精度 Tensor Cores。
方法论
本文扩展了 im2win (image to window) 范式,这是一种内存高效的卷积方法,通过转换输入张量来实现连续的内存访问和数据复用。作者在现代 GPU 架构(特别是 NVIDIA Ampere)上实现了该方法,并支持多精度:
- 数据变换: 与通过复制重叠窗口元素来创建高度冗余矩阵的
im2col 不同,im2win 将滑动窗口元素在不同通道间连续排列。这种变换通过消除对重叠元素的冗余存储(特别是在步长小于滤波器高度时),降低了内存使用量。
- 多精度实现:
- CUDA Cores: 实现全精度 (FP32) 内核。
- Tensor Cores: 使用 WMMA (Warp Matrix Multiply-Accumulate) API 实现半精度 (FP16) 内核。该设计利用 NHWC 布局来增强卷积窗口的内存访问连续性。
- 内核优化: 为了最大化 Tensor Cores 的硬件利用率,作者引入了几项特定的优化:
- 索引预计算 (Index Precomputation): 偏移量在主机端预先计算并存储在常量内存中,以避免冗余的运行时计算。
- 异步数据移动 (Asynchronous Data Movement): 利用 PTX 指令将数据移动与计算分离,通过将数据传输与处理过程重叠来隐藏延迟。
- 之字形访问 (Zig-Zag Access): 根据线程块奇偶性重新排列数据访问顺序,以减轻共享内存的银行冲突 (bank conflicts)。
- 双缓冲 (Double Buffering): 使用两个共享内存缓冲区和寄存器集来创建一个生产者-消费者流水线,允许在计算当前数据块时预取下一个数据块。
- 执行策略: 实现过程同时处理所有批次(全批次策略),以实现更高的并行度,这与 cuBLAS 和 cuDNN 常用的单批次或小批量策略形成对比。
主要贡献
本文做出了三个主要贡献:
- 多精度扩展: 将
im2win 范式扩展到支持 CUDA 核心的全精度执行和 Tensor Cores 的半精度执行。
- 优化的内核设计: 通过一系列针对现代 GPU 高性能执行量身定制的优化(索引预计算、之字形访问、异步数据移动和双缓冲)增强了
im2win。
- 全面的评估: 在十二个不同的 CNN 基准测试中进行了彻底的实验评估,将优化的
im2win 与 PyTorch 的 im2col (cuBLAS) 以及多种版本的 cuDNN (IPG, FT, WN) 在 CUDA 和 Tensor Cores 上进行了对比。同时还进行了消融研究,以量化单个优化技术的影响。
实验结果
评估是在配备 NVIDIA GeForce RTX 3090 (Ampere 架构) 的设备上进行的,涵盖了各种卷积核尺寸和输入维度的 12 个不同卷积基准测试 (cv1–cv12)。
- 性能 (TFLOPS):
- Tensor Cores:
im2win FP16 实现的 TFLOPS 最高可达其自身 CUDA 核心 (FP32) 实现的 2.7 倍。它比 cuDNN 的 Tensor Core 实现提升了 1.4 倍,比基于 im2col 的 cuBLAS 卷积提升了 6.4 倍。
- CUDA Cores: 在特定基准测试中,
im2win (FP32) 优于 cuDNN 的隐式 GEMM 和 Winograd 变体,相对于 cuDNN_WN 实现了高达 2.5 倍的 TFLOPS 提升。
- 内存效率:
im2win 显著降低了内存消耗。在 Tensor Cores 上,其使用的内存仅约为 cuDNN 的 53% 以及 im2col-cuBLAS 所需内存的 35%。
- 在 CUDA 核心上,它在大多数基准测试中始终比 cuDNN 变体使用更少的内存。
- 消融研究:
- 双缓冲 (Double Buffering) 被确定为最关键的优化,通过有效隐藏数据传输延迟提供了最大的性能增益。
- 异步数据移动 (Asynchronous Data Movement) 提供了第二显著的改进。
- 之字形访问 (Zig-Zag Access) 带来的收益最为微小,因为其他优化已经降低了发生银行冲突的概率。
重要性
本文将 im2win 确立为现代 GPU 架构下统一的高性能卷积框架。通过解决 im2col 的高内存开销和直接卷积的低局部性问题,同时利用 Tensor Cores 的混合精度能力,所提出的方法为深度学习工作负载提供了可扩展的解决方案。结果表明,im2win 可以作为现有先进标准(尤其是当内存带宽和 Tensor Cores 利用率是关键约束时)的强有力替代方案。这项工作证明,通过精心的算法设计结合硬件特定的优化(如异步数据移动和双缓冲),可以在成熟的标准(如 cuDNN 和 cuBLAS)之上获得实质性的性能和效率提升。
每周获取最佳 computer science 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。