PyTorch C++ 1.5.1在Ubuntu aarch64平台运行速度偏慢问题咨询

关于arm64版本和x64版本的优化程度差异 #

你使用的PyTorch 1.5.1是2020年发布的版本，该版本对arm64架构的优化程度远低于x86_64版本：

• x86_64版本在该迭代阶段已经积累了多年的优化沉淀，全量算子都适配了AVX、AVX2等x86专属指令集，核心算子还有大量手写汇编优化，性能表现已经非常成熟
• 同期arm64版本的适配还处于早期阶段，大部分算子仅实现了通用逻辑，未做NEON指令集的针对性优化，很多场景下会直接走性能较低的兜底实现，本身就和x64版本存在明显性能差距

关于数学库替换带来的性能影响 #

你提到的MKL替换为OpenBLAS+Eigen是导致性能差距的核心原因之一：

• MKL是英特尔针对x86架构深度定制的数学计算库，针对深度学习核心的矩阵乘法、卷积等算子做了极端优化，性能远高于通用开源数学库
• PyTorch 1.5.1 arm64版本默认绑定的OpenBLAS，在arm64架构下的矩阵运算性能仅为同算力x86平台MKL的30%~50%，刚好匹配你测试得到的1/3性能差。另外该版本尚未集成ARM官方的Compute Library（ACL），没有用到arm平台专属的数学库优化，进一步拉大了性能差距

优化建议 #

• 先检查当前PyTorch aarch64的编译参数，确认已开启-O3优化、NEON指令集支持，同时确认链接的是多线程版本的OpenBLAS，默认编译的单线程OpenBLAS性能会低40%以上
• 如果必须沿用PyTorch 1.5.1版本，可以手动对接ARM Compute Library作为算子后端，核心运算性能可以提升1倍以上
• 如果允许升级框架版本，建议升级到PyTorch 2.0及以上版本，该版本的arm64适配已经成熟，默认集成ACL后端，和同算力x64版本的性能差距可以缩小到20%以内

内容的提问来源于stack exchange，提问作者Gerald