TensorFlow CPU加速实战：从MKL到oneDNN的性能调优与部署实践

1. CPU跑TensorFlow慢问题不一定出在算力1.1 现象同样的模型在不同机器上跑出了两种结果不少同行应该都遇到过这个场景新开一台机器pip install tensorflow装完之后把之前写好的一套模型代码原封不动跑一遍发现 CPU 占用率确实拉上去了但一个 batch 的推理时间反而比旧机器还慢。我当初排查第一反应是数据管道、磁盘 IO 或者 Python 环境出了问题折腾了半天最后才发现根子不在代码上而在 TensorFlow 本身的 CPU 算子实现和这台机器的指令集不匹配。这里就要引出 Intel MKL 这个工具。MKL 的全称是 Math Kernel Library简单说是一套针对 Intel 处理器高度优化过的数学核心函数库。矩阵乘法、卷积、傅里叶变换这些深度学习里最吃计算量的操作MKL 都提供了比通用实现更快的内核。TensorFlow 在 CPU 上做推理和训练时可以把一部分算子从默认的后端切到 MKL 后端从而利用 AVX2、AVX-512 这些 SIMD 指令集以及 cache 友好的分块策略把 CPU 的算力彻底压出来。我自己实际接触下来这个工具的收益不是玄学。同一个 TextCNN 文本分类模型在一颗 i7-8700 的机器上默认 TensorFlow 跑 100 个 batch 大概 0.85 秒切到 Intel 优化版 TensorFlow 并做好线程配置之后压到了 0.55 秒左右。而在服务器端的一颗 EPYC 7302 上跑一个 Embedding Dense 的推荐排序模型开启 oneDNN 支持之后吞吐提升了大约 1.6 倍。当然这个提升幅度会因为算子类型、输入尺寸、机器本身有波动但大方向几乎没有悬念只要你在 CPU 上跑 TensorFlowMKL 这条优化路线值得花时间搞清楚。1.2 默认实现慢在哪儿从矩阵乘法说起要理解 MKL 为什么快得先明白 CPU 上的矩阵乘法难在哪里。平时我们写matmul(A, B)时脑子里想的是三重循环但真正的访存开销往往比计算开销更致命。CPU 从内存拿数据的速度远赶不上计算单元消耗数据的速度所以你让一个通用实现的 GEMM 去跑一个 512x512 的矩阵乘大部分时间其实都在等数据。通用实现通常假设数据是规规矩矩的连续行主序循环顺序也相对简单这样能正确处理任意形状的矩阵。但代价是 cache 命中率不稳定SIMD 向量化也不彻底。MKL 做的事情本质上就是三件第一把矩阵分块成能装进 L1/L2 cache 的小块第二根据 CPU 支持的指令集自动选择对应的 SIMD 内核比如 AVX2 一次能处理 8 个 floatAVX-512 一次能处理 16 个第三针对不同形状的矩阵走不同的分支策略小矩阵走 direct kernel大矩阵走拷贝分块后的 packed kernel。TensorFlow 里很多算子都依赖 Eigen 这个模板库做基础运算Eigen 本身也写得很好但它是通用库需要兼顾不同编译器、不同架构、不同操作系统。MKL 则是针对 Intel 架构反复调优过的专精方案所以在矩阵乘、卷积这类算子上的性能往往能拉开明显差距。1.3 MKL 不是万金油收益上限由算子构成决定这里得给大家先泼一盆冷水。MKL 加速并不是所有代码都能无脑快一倍。你如果模型里全是向量逐个元素的激活操作、Copy 操作、或者大量小 shape 的矩阵乘MKL 的优势会被布局转换和线程调度的开销抵消掉一部分。尤其是算子内部如果涉及 NHWC 和 block layout 之间的转换数据搬运的开销可能直接吃掉计算加速的红利。我自己测试时就有个印象很深的例子一个以 1x1 卷积为主的轻量型网络切到 MKL 之后并没有比默认 Eigen 快多少某些 batch 上甚至略慢。原因就是布局转换太频繁。反观包含大矩阵乘、3x3 卷积、较大 batch 的模型加速效果肉眼可见。所以后面所有调优实验我都建议先对自己的模型做一遍算子构成分析再决定要不要在 MKL 这条路上下重注。2. 从 MKL 到 oneDNN这个加速工具到底改了什么2.1 简单理清 MKL、MKL-DNN、oneDNN 的关系如果你自己去搜索相关材料大概率会被这几个名字绕晕Intel MKL、MKL-DNN、DNNL、oneDNN再加上现在经常看到的 Intel oneAPI。我最初也被搞糊涂过后来梳理清楚之后就简单了。Intel MKL 是历史最悠久、覆盖最广泛的数学库。TensorFlow 早年主要用它来加速矩阵乘和直卷积。后来 Intel 从 MKLDNN 衍生出专门面向深度神经网络的原语库也就是 MKL-DNN再后来又把它捐给了 Linux 基金会改名为 oneDNN。所以你在 TensorFlow 编译配置里看到的--configmkl、运行时看到的oneDNN、环境变量TF_ENABLE_ONEDNN_OPTS其实是一脉相承的东西。对使用者来说这些名字的演变带来的最大影响是不同版本的 TensorFlow启用 MKL 加速的方式和日志输出格式不完全一样。老版本可能认MKL_VERBOSE新版本更偏向DNNL_VERBOSE或ONEDNN_VERBOSE。后面我会单独拿出一个小节细说。2.2 官方包其实已经内置了一部分加速能力很多人不知道TensorFlow 的官方 pip 包在 Linux x86_64 平台上从 2.9 那个版本周期开始默认就编译了 oneDNN 支持。更早一点在 2.5 到 2.8 之间你还需要通过设置环境变量TF_ENABLE_ONEDNN_OPTS1来手动打开。这就引出一个很常见的误判某些博主或者社区帖子还在说“默认 TensorFlow 没有 MKL 支持必须自己编译安装”这在新版本里已经不完全准确了。如果你用的恰好是官方 Linux pip 包可以先不用重装在命令行里执行一下export TF_ENABLE_ONEDNN_OPTS1然后跑一个小卷积模型观察日志或者系统线程情况。当然2.9 之后的版本如果默认已经开启这个变量就不会生效需要换个方式验证这个我放在第四章详细讲。2.3 Intel 自己维护的 TensorFlow 包又是怎么回事除了官方 TensorFlowIntel 还维护了一套直接发布在 PyPI 上的优化版包名字就是intel-tensorflow。它主要有几个特点编译时针对 Intel CPU 开得更狠可能包含 AVX-512、VNNI 这些指令集oneDNN 相关的优化默认全部打开同时会绑定 Intel 自家的 LLVM 工具链来提升自动向量化效果。在环境允许的情况下换装这个包是最省事的 MKL 加速路线。pip install intel-tensorflow装完之后原本的 TensorFlow 代码基本不用改。但这里有个小坑intel-tensorflow的版本跟官方包不一定严格同步有些公司内部如果锁死了 TensorFlow 版本号升级前必须做一轮兼容性验证。3. 三种部署方式实测选对方案比猛调参数更重要3.1 最省事方案直接安装 intel-tensorflow如果你只想把已有项目快速切到 MKL 加速上我建议第一选择就是intel-tensorflow。步骤如下pip uninstall tensorflow pip install intel-tensorflow注意尽量在干净环境里操作。我之前在一个已经装过官方 TensorFlow 和 PyTorch 的 conda 环境里直接覆盖结果遇到了 OpenMP 库冲突Python 进程一启动就崩。后文踩坑部分会专门讲。装完之后运行一个最简单的验证脚本import tensorflow as tf print(tf.__version__) print(tf.sysconfig.get_build_info())在返回信息里你通常能看到cpu_feature_guard相关日志里面有Intel(R) MKL-DNN或oneDNN字样。如果没有就说明当前环境没有正确启用需要继续排查。这个方案的优点非常明显不需要编译不需要配置 Bazel半小时内能完成验证。缺点也明显就是没法针对你的特定 CPU 型号做更细的指令集定制而且它绑定的 Numpy、absl 等依赖版本可能和项目里其他模块冲突。3.2 兼顾可控性和性能官方 TensorFlow oneDNN 开关第二种路线是继续用官方包但通过环境变量和 TensorFlow API 把 oneDNN 打开同时对线程做配置。适合同事之间协作的团队项目因为这套方案不改变基本包依赖其他人拉代码下来也能跑。需要关注的环境变量有这么几个export TF_ENABLE_ONEDNN_OPTS1 export OMP_NUM_THREADS16 export KMP_BLOCKTIME1 export KMP_AFFINITYgranularityfine,compact,1,0代码层面如果你用的是 TensorFlow 2.x可以这样设置import tensorflow as tf tf.config.threading.set_intra_op_parallelism_threads(16) tf.config.threading.set_inter_op_parallelism_threads(2)这个配置的意思是单个算子的内部并行度用 16 个线程算子之间用 2 个线程并行。这样设置是为了避免所有算子同时抢线程导致上下文切换开销过大。我实测在 16 核以上的机器上intra_op高一点、inter_op保持 1 到 4整体吞吐通常更好。3.3 追求极限从源码编译 TensorFlow 并开启 MKL 指令集如果你负责的是需要长期运行、对单机性能极敏感的推理服务那么可以考虑源码编译。这个过程会比前面两种方式痛苦得多但能做到按需裁剪。源码编译最核心的步骤是从 GitHub 拉取 TensorFlow 源码然后配置 Bazel 编译参数。大致流程如下git clone https://github.com/tensorflow/tensorflow cd tensorflow ./configure配置过程中会问你是否启用 MKL选择y即可。接着执行编译bazel build --configmkl --configopt //tensorflow/tools/pip_package:build_pip_package如果你确认当前 CPU 支持 AVX-512还可以在编译参数里加上--copt-mavx512f --copt-mavx512bw --copt-mavx512vl --copt-mavx512dq但这里要提醒一句这些-mavx512参数会让编译出来的包只能在同样支持 AVX-512 的 CPU 上运行。如果服务器之后要迁移到较旧的机型这个二进制会直接 Illegal instruction。所以不建议在通用发布环境里开满所有指令集。源码编译的另一个痛点是时间。我曾在 32 核机器上编译 TensorFlow 2.8大概花了一个半小时期间内存占用稳定在 16 GB 以上。如果机器资源不够建议给 Bazel 加上--local_ram_resources8192这类参数限制它的资源占用免得编译过程中 OOM。3.4 容器化部署时的特殊考虑用 Docker 部署 TensorFlow 推理服务时MKL 加速的收益依然成立但要注意线程亲和性在容器里容易被忽视。容器默认情况下OMP_NUM_THREADS如果设置得太高线程会被调度到不同 NUMA 节点的核心上跨节点访存的代价会明显抵消加速效果。比较好的做法是启动容器时加docker run --cpuset-cpus0-15 --cpuset-mems0 your_image同时设置OMP_NUM_THREADS16。这样进程基本能扎在同一个 NUMA 域里内存访问更稳定。没有--cpuset-mems时最好也确认一下容器 Runtime 对 NUMA 的感知能力不然你调了半天性能参数可能瓶颈根本不在计算而在内存距离。4. 安装后先别急用 MKL_VERBOSE 和基准测试确认加速生效4.1 第一锤看运行日志里的 oneDNN/MKL 标记很多人在切完环境后跑了一通模型觉得“好像快了一点”但说不清到底有没有用上 MKL。我建议安装后的第一件事不是跑完整训练而是先用一段小脚本确认后端真的切过去了。拿 TensorFlow 2.x 举例开启 verbose 日志export MKL_VERBOSE1然后运行import tensorflow as tf x tf.random.normal([256, 256]) y tf.random.normal([256, 256]) for _ in range(5): tf.matmul(x, y)在较新的 oneDNN 版本里可能看不到带MKL_VERBOSE的输出这时试试export DNNL_VERBOSE1如果看到了类似oneDNN: verbose ... gemm ... jit:avx512或者DNNL_VERBOSE ... convolution这样的行就能确认 MKL/oneDNN 内核确实被调用了。如果什么都没有说明 TensorFlow 在编译或运行时压根没有启用这条后端的路径后面的调优全都无从谈起。4.2 第二锤用真实模型做基准而不是跑 hello world日志只能证明后端起没起不能直接说明收益有多大。想量化收益一定要拿自己业务里真实会跑的模型测。原因在于不同算子比例下加速差异非常大。我只拿一个随机矩阵乘来测试没什么意义。建议写一个简单的基准脚本import time import numpy as np import tensorflow as tf # 模拟一个固定 shape 的输入 model tf.keras.Sequential([ tf.keras.layers.Dense(512, activationrelu), tf.keras.layers.Dense(512, activationrelu), ]) x tf.random.normal([64, 256]) # 预热 for _ in range(10): model(x) start time.perf_counter() for _ in range(100): model(x) end time.perf_counter() print(favg time per inference: {(end - start) / 100 * 1000:.2f} ms)在切换 MKL 开关前后分别跑一遍同一份脚本记录平均耗时。我通常会把结果跟机器信息一起记到笔记里因为 CPU 频率波动、虚拟化环境里的邻居干扰都会影响结果单跑一次不能作数。为了对比更可靠至少跑三轮每轮之间休息几分钟取中位数。如果三轮结果波动超过 15%先别急着下结论查一下机器是不是被其他进程占用了。4.3 第三锤看线程数加载曲线确认并行度真实生效日志和耗时都验证通过之后还有一步值得做跑到一个计算密集的模型里去开另一个终端看 CPU 占用率分布。top -u $USER或者用htop看每个核的情况。如果OMP_NUM_THREADS16那计算密集型算子的 CPU 占用应该大致在 1600% 附近用 top 显示时是 1600% 的进程。如果看到系统 CPU 占用高但进程 CPU 占用低那问题通常不在 MKL而在数据加载或者 Python 侧的锁竞争。我自己曾经踩过这种坑MKL 日志显示内核都起来了但进程 CPU 占用就是上不去最后发现是前一个进程残留的 CPU affinity 设置把进程限制在了两个核上。所以验证时别只看 MKL 日志还要确认操作系统层面真的把线程铺开了。5. 线程与内存参数调优把CPU的最后一点性能抠出来5.1 intra_op 和 inter_op 的配合逻辑TensorFlow 在 CPU 上有两层并行度intra_op_parallelism_threads表示单个算子内部的并行线程数inter_op_parallelism_threads表示同时执行多少个相互独立的算子。拿一个 Dense 层举例矩阵乘阶段由intra_op控制如果你的计算图里有多个分支可以同时跑那么多个分支之间由inter_op控制。这两者的最优组合并不是越大越好。我之前在一台两路服务器上把intra_op调到 32、inter_op调到 32结果性能反而下降了三成。原因很简单算子内部和算子之间的线程同时竞争 CPU 资源线程频繁切换缓存反而被打乱了。一个相对稳妥的启动配置是import tensorflow as tf cpu_count os.cpu_count() intra_op min(cpu_count, 16) inter_op 2 tf.config.threading.set_intra_op_parallelism_threads(intra_op) tf.config.threading.set_inter_op_parallelism_threads(inter_op)这个配置不适合所有场景但适合大多数中小型在线推理任务。如果你跑的是大批量离线推理可以适当把inter_op调低到 1让所有线程集中精力处理单个算子减少线程间同步。5.2 KMP_BLOCKTIME 到底该设多少KMP_BLOCKTIME是 Intel 的 OpenMP 运行时参数含义是线程在等不到新任务时主动休眠之前的时间单位是毫秒。默认值往往是 200这在以计算为主的深度学习任务里偏大因为线程会白白等太久。一般做推理服务时我会把它调到 1export KMP_BLOCKTIME1这样线程执行完当前任务后很快进入休眠等待把 CPU 空出来给别的线程或进程。对于单进程内部的算子切换来说这个值通常能降低延迟抖动。但如果你的进程里反复有大量小算子启动过小的KMP_BLOCKTIME可能导致线程反复睡眠唤醒反而增加开销。所以这个值最好是实测着调先 1再 10对比 1000 次稳定延迟。5.3 超线程和 NUMA 场景下的线程绑定策略现在的服务器 CPU超线程往往让操作系统识别到两倍物理核数量的逻辑核。OMP_NUM_THREADS设置成跟逻辑核数一样时不一定最快。我测试过一个 24 物理核的机器OMP_NUM_THREADS24和OMP_NUM_THREADS48相比前者的矩阵乘吞吐反而高出约 12%。原因是超线程共享物理核的执行单元计算密集型负载并不能获得两倍资源。如果你的机器是单路 CPU 或者虚拟化实例绑核不是必须的但多路 CPU 一定要关注 NUMA。最省事的做法是直接用numactl控制numactl --cpunodebind0 --membind0 python train.py这样进程只会在 Node 0 的 CPU 上运行内存也优先分配在 Node 0。如果模型小、数据量也小这种做法可以避免跨 NUMA 访问效果显著。5.4 内存带宽可能就是你的隐藏瓶颈MKL 把计算优化得很好之后很多场景的瓶颈会从 CPU 计算转移到内存带宽。比如大 batch 的卷积操作特征图数据量巨大一遍遍从主存搬到 cache再算完写回。这时候你再调线程参数已经没有太大收益需要考虑的是数据布局和内存分配。我遇到过一个实际问题同样的 ResNet-50 推理batch size 从 1 改成 16 之后耗时不是线性增长而是在某一个点突然暴涨。分析之后发现是内存带宽已经饱和。解决办法是把输入图缩小、走 TFRecord 流水线分 batch 喂、或者干脆降低单次请求的 batch分几个小 batch 异步处理反而整体吞吐更高。这类问题不会在 MKL 日志里暴露需要靠 profiling 工具比如 Intel 的 VTune或者至少盯紧perf stat里的LLC-load-misses。6. 最常遇到的五个坑以及完整的排查思路6.1 报 “AVX AVX2 not used”是不是一定要换二进制启动 TensorFlow 时经常看到这行警告Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX AVX2。意思是当前 Python 包里的二进制没有用上你这颗 CPU 的 AVX/AVX2 指令集但 TensorFlow 仍然能跑只是性能没有到最优。这种提示出现时最优先的应对方式是装intel-tensorflow或者自己编译。但如果你只是想快速关闭日志倒是可以设置环境变量export TF_CPP_MIN_LOG_LEVEL2我个人的建议是不要只关日志了事因为这本质上是在提示你当前性能有很大缺口。一次矩阵乘从 SSE 换成 AVX2性能可能差在一两倍以上。如果你部署的是长期运行的线上服务这值得认真升级。6.2 Python 直接崩libiomp5md.dll / libiomp5.so 重复加载这个坑在同时装有 PyTorch 和 TensorFlow 的环境里太常见了。报错信息大致是OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized.根因是 PyTorch 和 Intel 优化版 TensorFlow 各自带了不同版本的 OpenMP 运行时两者同时加载时发生冲突。网上流传的应急办法是设置export KMP_DUPLICATE_LIB_OKTRUE这个变量能绕过错但它掩盖了真正的问题有可能造成线程池行为不可预期。更干净的办法是给两个深度学习框架分别建虚拟环境彻底隔离原生库。实在要共存可以尝试用conda install nomkl强制让其中一个框架使用 GNU 的libgomp减少 Intel OpenMP 的重复加载概率。6.3 编译 TensorFlow 时中途失败先看这三样源码编译的失败原因多到可以单独写一篇但我打过这么多次交到绝大多数问题集中在三处。第一处是磁盘空间。Bazel 的 cache 会缓存大量中间产物TensorFlow 整个编译过程轻轻松松吃掉 30 GB 以上。建议编译前用df -h检查一下给/root/.cache/bazel所在分区留足 50 GB。第二处是内存。用虚拟机编译时内存少于 8 GB 基本没戏干脆用--local_ram_resources限制编译并行度哪怕慢一点也别 OOM。第三处是 JDK 版本。TensorFlow 构建工具链对 JDK 版本有要求尤其是比较老的版本装一个不匹配的 JDK 会直接在 configure 阶段报错。如果你不想折腾这些又必须源码编译可以先用官方提供的tensorflow/tensorflow:develDocker 镜像作为基础环境里面的编译依赖都已经配好能省掉一大半前置问题。6.4 某些算子反而变慢布局转换的开销MKL 在后端优化时经常会把 TensorFlow 默认的 NHWC 数据布局转换成内部最擅长的 blocked layout。这个转换在算子少、网络浅的模型里可能在每个算子前后都做一次导致内存复制量比计算量还大。如果你发现启用 MKL 后某些层耗时反而增加可以先打开 oneDNN 的 verbose 日志看是不是有大量reorder操作。如果reorder的耗时占比很高尝试把输入数据的 format 在代码里固定下来或者在tf.keras.layers外边多用tf.ensure_shape明确 shape 信息减少不必要的布局推导。最直接的办法是拿 Intel 官方的 Model Zoo 里的脚本对比一下如果同样是 ResNet-50官方脚本快很多那说明问题出在你自己的数据加载或布局构造上。6.5 多进程服务里的线程池相互踩踏最后一种坑很隐蔽发生在你用多进程部署推理服务时。比如你用multiprocessing起了四个 worker每个 worker 在初始化 TensorFlow 时都直接读取了默认的OMP_NUM_THREADS结果四个进程各自开了 16 个线程一台 32 核机器瞬间被 64 个线程淹没。这种情况下单看单个进程的耗时可能没什么异常但从整体看吞吐上不去延迟反而变高。解决办法是在每个 worker 创建 TensorFlow 会话之前单独设置OMP_NUM_THREADS为总物理核数除以进程数。更细的做法是配合kmp_set_thread_affinity_mask之类的接口把不同进程的线程绑定到不同核上。我在做线上推理服务时专门写过一个环境变量分发逻辑在启动 worker 时给每个子进程注入独立的线程数配置这才彻底解决了线程踩踏问题。最后再分享一个我自己的经验每次换 CPU 加速相关方案时我都坚持固定三个变量模型、batch size、数据集。只改 TensorFlow 的安装方式或线程配置这样对比出来的数据才有说服力。否则今天你顺便升级了 Numpy明天换了机器后天又把 batch 调大了最后根本说不清快在哪、慢在哪。MKL 加速这条路其实不难走难的是每一步都被验证过、被记录过而不是靠感觉调参。希望这篇复盘能帮你少走我走过的那些弯路。