中文

基于 SIMD 架构的三种张量布局的高性能 Im2win 与直接卷积

机器学习 2024-08-02 v1 人工智能 神经与进化计算

摘要

卷积是深度神经网络的核心组件,计算密集且耗时。张量数据布局对卷积操作的内存访问和计算效率有显著影响。然而,目前仍缺乏对 SIMD 架构上数据布局在卷积方法方面的全面性能表征。本文为 im2win 卷积提出了三种新型数据布局:NHWC、CHWN 和 CHWN8,并引入了一套针对直接卷积和 im2win 卷积的通用优化技术。我们将优化后的 im2win 卷积与直接卷积以及 PyTorch 的 im2col 卷积在上述布局上进行比较。实验表明,采用新 NHWC 布局的 im2win 卷积相比 NCHW 布局实现了高达 355% 的性能加速。我们的优化还显著提升了 im2win 和直接卷积的性能。优化后的 im2win 和直接卷积分别达到了机器理论峰值的 95% 和 94%。

关键词

引用

@article{arxiv.2408.00278,
  title  = {High Performance Im2win and Direct Convolutions using Three Tensor Layouts on SIMD Architectures},
  author = {Xiang Fu and Xinpeng Zhang and Jixiang Ma and Peng Zhao and Shuai Lu and Xu T. Liu},
  journal= {arXiv preprint arXiv:2408.00278},
  year   = {2024}
}