如何限制TensorFlow的CPU内存占用？TensorFlow新手求助

嘿，作为曾经踩过TensorFlow CPU内存坑的过来人，给你梳理下问题和解决办法：

首先得明确：你脚本里的tf_config.gpu_options.allow_growth只对GPU内存分配生效，完全不影响CPU内存，这是新手常犯的误区～

TensorFlow默认会尽可能占用所有可用的CPU资源（包括内存和线程）来优化性能，但对于小脚本来说完全没必要，我们可以通过配置tf.ConfigProto来限制它：

1. 调整CPU线程和核心占用 #

修改你的配置部分，添加CPU相关的限制参数：

tf_config = tf.ConfigProto(
    # 限制使用的CPU核心数，根据你的需求改数值（比如1或2）
    device_count={'CPU': 1},
    # 控制不同操作之间的并行线程数
    inter_op_parallelism_threads=1,
    # 控制单个操作内部的并行线程数
    intra_op_parallelism_threads=1,
    # 原有的GPU配置保留即可
    gpu_options=tf.GPUOptions(allow_growth=False)
)

这些参数会让TensorFlow不再抢占全部CPU资源，内存占用会明显降低。

2. 关闭不必要的优化（可选） #

如果内存占用还是超预期，可以关闭XLA即时编译优化，进一步减少内存开销：

tf_config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.OFF

修改后的完整脚本 #

# -*- coding: utf-8 -*-
import time
import tensorflow as tf

tf_config = tf.ConfigProto(
    device_count={'CPU': 1},
    inter_op_parallelism_threads=1,
    intra_op_parallelism_threads=1,
    gpu_options=tf.GPUOptions(allow_growth=False)
)
tf_config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.OFF

with tf.Session(config=tf_config) as sess:
    print('Listening.....')
    time.sleep(100)

另外补充下：导入TensorFlow时占用约100MB内存是正常的，这是初始化核心组件的开销；而Session启动后内存飙升，本质是TensorFlow默认创建了大量线程池并预分配内存，上面的配置就是针对性地限制这些行为。

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