TensorFlow加载模型至GPU后CPU内存仍被占用，如何释放？

现象是否正常？ #

这种情况不算完全正常：TensorFlow加载模型时，CPU端确实会保留少量元数据、计算图结构等必要内存，但如果CPU内存占用过高且全程不释放，就属于需要优化的异常情况。

解决办法 #

1. 限制TensorFlow的CPU内存使用 #

直接给CPU设置内存上限，避免无限制占用：

import tensorflow as tf

# 配置CPU内存限制，示例为2GB
cpu_devices = tf.config.list_physical_devices('CPU')
if cpu_devices:
    tf.config.set_logical_device_configuration(
        cpu_devices[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=2048)]
    )

with tf.device("/GPU:0"):
    my_model = tf.keras.models.load_model(filepath=filepath)

2. 加载后主动清理CPU冗余对象 #

加载完成后，手动删除临时变量并强制垃圾回收：

import tensorflow as tf
import gc

with tf.device("/GPU:0"):
    my_model = tf.keras.models.load_model(filepath=filepath)

# 清理加载过程中产生的临时TensorFlow对象（若无需再调用TF API可删除tf引用）
del tf
gc.collect()

3. 改用底层SavedModel加载方式 #

用tf.saved_model.load替代Keras的加载方法，减少CPU端的封装开销：

import tensorflow as tf

with tf.device("/GPU:0"):
    my_model = tf.saved_model.load(filepath)

4. 检查并迁移CPU绑定的层 #

部分自定义层或旧层可能默认绑定CPU，导致权重留在CPU内存。可以遍历检查并迁移：

# 查看每层的设备分配
for layer in my_model.layers:
    print(f"Layer {layer.name}: {layer._device}")

# 将所有层迁移到GPU
for layer in my_model.layers:
    layer._device = "/GPU:0"

补充说明 #

• TensorFlow的tf.device只是指定优先加载设备，框架会在CPU保留少量管理内存，这部分无法完全释放，只要占用量远低于GPU内存就无需过度纠结。
• 生产环境中如果代码改动受限，可先初始化GPU环境再加载模型，减少CPU预分配：

import tensorflow as tf

# 开启GPU内存动态增长，避免预占过多
gpu_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpu_devices[0], True)

with tf.device("/GPU:0"):
    my_model = tf.keras.models.load_model(filepath=filepath)

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