趋近智
内存物理上是一维的。深度学习框架将张量表示为具有 等形状的多维对象,但硬件看到的是字节的平坦序列。逻辑张量索引与物理内存地址之间的映射决定了内存布局。编译器选择内存布局是图级优化中非常重要的决定之一,因为它直接影响数据局部性、缓存使用率以及专用硬件指令的使用能力。
张量布局与硬件固有要求不匹配会导致分散的内存访问模式。这会导致高缓存未命中率,并阻止使用像Tensor Cores或AVX-512这样的高吞吐量指令。布局转换是编译器的一个阶段,它负责重写图以确保数据按计算引擎期望的顺序存储在内存中。
操作效率与内存访问步长相关。考虑一个标准的2D卷积。该操作涉及输入通道和滤波器权重之间的点积。
如果我们以 格式(批次、通道、高度、宽度)存储数据,则最内层维度是宽度。图像中空间上相邻(宽度方向)的值在内存中也是相邻的。然而,在不同通道但相同空间位置的值之间,相隔 个元素。
与此相对,考虑 格式(批次、高度、宽度、通道)。在这种格式下,最内层维度是通道。对应于所有通道中相同像素的值被连续存储。由于现代卷积实现通常在通道维度上进行归约(沿深度方向累积输入), 允许硬件在单次事务中加载密集向量的通道数据。这显著提高了GPU上的内存合并效率,并使CPU上的向量化成为可能。
我们可以将 布局中4D张量索引 的内存地址 定义为:
对于 布局,映射变为:
当编译器将图转换为特定目标时,它会向后端查询首选布局。对于使用Tensor Cores的NVIDIA GPU(通过cuDNN或Cutlass),首选布局几乎总是 (通常称为通道在后)。对于依赖SIMD指令的CPU,最佳布局通常是分块格式,它创建的数据块大小刚好适合向量寄存器。
对于CPU目标,像 或 这样的标准布局通常不足以最大化算术强度。为了充分使用SIMD单元(例如AVX-512或ARM Neon),编译器经常使用布局分块(也称为平铺或打包)。
分块是指将一个维度拆分为一个外维度和一个固定大小 的内维度。例如,我们可以将通道维度 转换为 。这将4D张量 转换为5D张量 ,其中小写 表示一个连续存储的通道块(例如8或16)。
如果我们选择块大小为16,布局变为 。内存地址计算公式变为:
这种结构确保当CPU处理一个空间像素时,它在一条指令中将精确的16个通道加载到512位向量寄存器中。这消除了对收集/分散指令的需求,并确保数据与FMA(乘加运算)单元完美对齐。
下图说明了布局选择对通用卷积操作内存吞吐量的性能影响。
不同张量布局的有效内存带宽利用率比较。分块和通道在后布局通过最小化缓存行逐出,在现代密集架构上显著优于平面 格式。
改变单个操作符的布局很简单;但处理对整个图的影响则很复杂。如果卷积从 转换为 ,其输入必须进行转置。如果后续操作(例如,批归一化或ReLU)期望 ,则输出必须转置回来。
在每个节点前后插入显式的 Transpose 操作会违背优化的目的。内存带宽成本高昂,转置大型张量是一种内存密集型操作,会消耗大量时间和能量。为了解决这个问题,编译器实现了一个名为布局传播的阶段。
在这个阶段,编译器根据硬件目标,为特定的“锚定”操作符(通常是卷积或矩阵乘法)分配首选布局。然后它遍历图以将此布局要求传播给相邻操作符。诸如ReLU、Add或Sigmoid等逐元素操作与布局无关;它们可以像处理 数据一样轻松地处理 数据,而无需改变其数学定义。
编译器通过这些与布局无关的节点推动布局转换,直到遇到必须改变布局的边界(例如,一个Reshape操作或外部输出)。这使得编译器可以将整个子图转换为目标布局,从而将图边界处必要的转置次数降至最低。
下图描绘了布局传播阶段中子图的转换过程。
布局传播的可视化。编译器识别出Conv2D需要NHWC布局。编译器没有只在Conv2D周围进行转置,而是将NHWC布局通过与布局无关的ReLU传播出去,将耗时的逆转换移至序列的末尾。
实现布局转换需要细致处理操作符属性。当布局改变时,编译器不仅要更新输入张量,还要更新操作符的静态属性。
stride(步长)、padding(填充)或 dilation(膨胀)。如果布局旋转,这些属性必须重新索引。例如,[2, 2] 的步长通常适用于 。如果布局改变,使得空间维度从索引2和3移动到索引1和2,则步长属性向量必须进行置换以匹配。目标是确保操作的语义含义保持一致,同时物理执行适应硬件的优势。通过自动化此过程,编译器将模型定义与硬件实现解耦,使相同的高级Keras或PyTorch代码能够在移动CPU(使用 )和数据中心GPU(使用 )上高效运行。
这部分内容有帮助吗?
© 2026 ApX Machine Learning用心打造