如何在服务器上限制TensorFlow 2与Keras的CPU核心使用数

1. 控制底层计算库的线程数（核心步骤） #

TensorFlow在Linux环境下会依赖OpenMP、MKL等底层计算库，这些库默认会占用所有可用CPU核心。必须在导入TensorFlow、numpy前设置以下环境变量：

import os
# 统一设置各计算库的线程数为5，可根据需求调整
os.environ["OMP_NUM_THREADS"] = "5"
os.environ["OPENBLAS_NUM_THREADS"] = "5"
os.environ["MKL_NUM_THREADS"] = "5"
os.environ["NUMEXPR_NUM_THREADS"] = "5"

2. 适配TF2.x的线程配置 #

你之前使用的TF1.x兼容配置在TF2.10中可能不生效，改用TF2.x原生API配置，且要在创建模型前执行：

import tensorflow as tf
# 单操作内部并行线程数
tf.config.threading.set_intra_op_parallelism_threads(5)
# 多操作间并行线程数
tf.config.threading.set_inter_op_parallelism_threads(2)
# 关闭自动并行优化，避免额外线程占用
tf.config.optimizer.set_jit(False)

3. 调整joblib的运行方式 #

即使设置n_jobs=1，multiprocessing后端仍可能出现线程泄漏。建议：

• 改用backend="threading"：

  from joblib import Parallel, delayed
  Parallel(n_jobs=1, backend="threading")(delayed(run_net)(param) for param in params)
  

• 或者直接去掉joblib，用纯循环执行：

  for param in params:
      run_net(param)
  

4. 系统层面强制限制（兜底方案） #

如果以上方法无效，用Linux系统工具taskset绑定进程到指定核心：

• 启动脚本时直接指定核心（例如绑定0-4号核心，共5个）：

  taskset -c 0-4 python your_script.py
  

• 运行中修改已启动进程的核心绑定：

  # 先获取进程PID
  ps aux | grep your_script.py
  # 绑定核心
  taskset -cp 0-4 <进程PID>
  

验证方法 #

运行脚本后，可通过以下命令确认核心占用情况：

# 查看进程的线程数量
ps -T -p <进程PID>
# 查看进程绑定的核心列表
taskset -p <进程PID>

内容的提问来源于stack exchange，提问作者Lutz Köhler