Keras/TensorFlow限制CPU核心数失效：intra_op_parallelism_threads不生效

我之前在共享CPU集群上也碰到过一模一样的问题——Keras/TensorFlow根本不听代码里的线程限制，直接把所有核心占满，搞得管理员来找我麻烦😅。结合我的踩坑经验，给你几个肯定能解决的方案：

1. 先设置环境变量，这是集群环境下最靠谱的方式 #

TensorFlow会优先读取系统环境变量来配置线程，代码里的设置反而可能被覆盖。在导入任何TensorFlow/Keras模块之前，先设置这些变量：

import os
# 强制限制各类线性代数库的线程数
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["OPENBLAS_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"
os.environ["VECLIB_MAXIMUM_THREADS"] = "2"
os.environ["NUMEXPR_NUM_THREADS"] = "2"

如果是用bash脚本提交作业，也可以在脚本开头加：

export OMP_NUM_THREADS=2

2. 确认Keras版本，选对配置方式 #

现在很多人混用keras和tf.keras，两者的session配置逻辑不一样：

• 如果你用的是tf.keras（推荐）：直接用TensorFlow原生的线程配置API，而且要在创建模型前执行：

  import tensorflow as tf
  # 先配置线程
  tf.config.threading.set_intra_op_parallelism_threads(2)
  tf.config.threading.set_inter_op_parallelism_threads(2)
  tf.config.set_soft_device_placement(True)
  
  # 再导入模型并创建
  from tensorflow.keras.applications.vgg16 import VGG16
  model = VGG16(weights=None, input_shape=(224,224,3))
  

• 如果你用的是独立的Keras库：需要手动设置Session，同样要在模型创建前：

  from keras.applications.vgg16 import VGG16
  from keras import backend as K
  import tensorflow as tf
  
  conf = tf.ConfigProto(
      device_count={'CPU': 1},
      intra_op_parallelism_threads=2,
      inter_op_parallelism_threads=2
  )
  session = tf.Session(config=conf)
  K.set_session(session)
  
  # 之后再创建模型
  model = VGG16(weights=None, input_shape=(224,224,3))
  

3. 关键！配置必须在模型初始化之前完成 #

我之前踩过的大坑：先创建了模型再去设置session/线程，这时候TensorFlow已经用默认的多线程初始化了模型，后续的配置完全无效。一定要把所有线程/环境变量配置放在model = VGG16(...)之前！

4. 配合集群调度系统的限制 #

如果你的集群用SLURM这类调度工具，提交作业时要明确指定CPU核心数，比如：

sbatch --cpus-per-task=2 your_script.py

这时候代码里的线程数要和--cpus-per-task的数值一致，不然系统层面的限制会覆盖你的代码配置。

最后测试一下，运行代码后用top或者htop看进程的CPU使用率，应该只会占用你设置的核心数了。

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