如何解决Python运行TensorFlow时CPU占用率过高的问题？

在TensorFlow性能调优中，线程配置是个常见但容易踩坑的环节。很多人发现set_intra_op_parallelism_threads没生效，其实原因很简单——设置时机不对。下面我们逐一拆解几个关键问题。

tf.config.threading.set_intra_op_parallelism_threads没生效？检查是否晚于TF初始化

这个函数有个硬性要求：必须在import tensorflow as tf之后、任何TensorFlow操作执行前调用。哪怕你只调了一次tf.constant(1)或定义了@tf.function，再调它就完全失效，而且连个错误提示都没有。

很多人犯的错误是，先import tensorflow as tf，接着model = tf.keras.Sequential(...)（触发内部图构建），最后才写tf.config.threading.set_intra_op_parallelism_threads(1)，结果就是无效。

正确的做法是，在import之后立即设置线程数：

import tensorflow as tf
tf.config.threading.set_intra_op_parallelism_threads(1)
tf.config.threading.set_inter_op_parallelism_threads(1)
# 此时才能创建张量、模型、@tf.function等
x = tf.random.normal((1000, 1000))

为什么设了线程数还是CPU打满？优先级被环境变量覆盖了

环境变量TF_NUM_INTRAOP_THREADS和TF_NUM_INTEROP_THREADS的优先级高于代码设置。Dockerfile、~/.bashrc、Slurm job script或PyCharm运行配置里残留的export都可能导致代码里的set_*被静默覆盖。

验证方式：启动Python后立即运行import os; print(os.environ.get("TF_NUM_INTRAOP_THREADS"))，再运行print(tf.config.threading.get_intra_op_parallelism_threads())，两者不一致，说明环境变量起了作用。

推荐做法：统一用环境变量启动，例如：

TF_NUM_INTRAOP_THREADS=2 TF_NUM_INTEROP_THREADS=2 python train.py

intra_op 和 inter_op 该设多少？别盲目填物理核心数

这两个参数正交作用，总并发线程数 ≈ intra_op × inter_op（实际受OS调度影响，非严格相乘）。设得过高反而因上下文切换拖慢速度。

• intra_op_parallelism_threads控制单个算子（如tf.matmul）内部用几个线程——设为1最稳妥，适合调试或避免调度开销
• inter_op_parallelism_threads控制多少个独立op可并发调度——在32核机器上设64不会提速，反而更慢；设2～4更适合混部或容器环境
• 设为0或负数会直接抛ValueError: Number of threads must be positive

CPU高但GPU利用率低？问题不在TF线程数，而在数据流水线

如果nvidia-smi显示GPU利用率长期低于30%，说明瓶颈在CPU数据供给，不是TF线程太多——此时调tf.config.threading没用，该优化tf.data。

• 确保用了.map(..., num_parallel_calls=tf.data.AUTOTUNE)
• 加.prefetch(tf.data.AUTOTUNE)让数据预加载与GPU计算重叠
• 避免在.map()中混用NumPy或PIL——必须用纯TensorFlow函数，否则无法并行

真正容易被忽略的是：线程限制必须在所有TF初始化动作之前完成，包括tf.distribute.Strategy实例化、甚至tf.keras.Model构造——它们内部都会触发eager执行或图准备，一旦发生，就锁死了线程配置。