如何禁用PyTorch C++ CPU端多线程？已配置线程参数仍多核心满载

问题根因 #

你遇到的16核心满负载问题，核心原因是PyTorch CPU算子依赖的底层数学计算库线程配置独立，不受PyTorch自身线程接口控制，或者是配置调用顺序错误导致设置失效。

必须补充的配置项 #

• 调整配置调用顺序：所有线程配置代码必须放在任意PyTorch相关操作（包括张量创建、模型加载、算子初始化）之前执行，PyTorch初始化完成后再修改线程配置不会生效。参考代码结构如下：

int main() {
    // 线程配置放在main函数最开头，优先执行
    at::set_num_interop_threads(1);
    at::set_num_threads(1);
    torch::set_num_threads(1);
    omp_set_num_threads(1);
    omp_set_dynamic(0);
    omp_set_nested(0);

    // 后续再执行模型加载、推理等业务逻辑
    // ...
}

• 补充底层依赖的线程环境变量：
  官方预编译的libtorch默认绑定MKL数学库和oneDNN算子加速框架，这两个组件的线程配置不会继承PyTorch的设置，需要单独配置对应的环境变量。如果用到TBB并行框架也需要同步配置。
  建议在程序启动前通过系统环境变量统一配置所有线程参数，避免代码内修改不生效的问题，启动参考命令：

# 配置所有相关依赖的线程数为1
export MKL_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export OMP_NUM_THREADS=1
export DNNL_NUM_THREADS=1
export TBB_NUM_THREADS=1
# 启动你的程序
./你的程序可执行文件路径

• （可选）操作系统层面绑定核心：如果需要彻底限制CPU核心使用，可以在Linux环境下通过pthread_setaffinity_np接口将进程绑定到单个CPU核心，从系统层面避免线程跨核心调度。

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