Ubuntu 16.4上训练TensorFlow模型时内存占用过高问题求解

我之前在Ubuntu 16.04上跑TensorFlow CPU训练的时候，碰到过和你一模一样的糟心问题——8G内存直接被吃满，OOM杀手立刻找上门，禁了OOM就直接Segfault。试了好几种方法，终于找到几个靠谱的解决方案，分享给你：

1. 给TensorFlow CPU模式加内存限制 #

别光盯着GPU参数，TensorFlow在CPU模式下默认会疯抢所有可用内存，你可以直接在代码里限制它的线程数，间接把内存占用压下来：

import tensorflow as tf

# 根据你的CPU核心数调整，比如4核机器设为2-3就行，别拉满
tf.config.threading.set_intra_op_parallelism_threads(2)
tf.config.threading.set_inter_op_parallelism_threads(2)

# 要是你用的是TensorFlow 1.x版本，换成下面的配置
# config = tf.ConfigProto(intra_op_parallelism_threads=2, inter_op_parallelism_threads=2)
# session = tf.Session(config=config)

这个方法能直接限制TF的计算并发度，避免它一次性占用过多内存。

2. 别一次性加载全量数据，改成分批读取 #

很多时候内存暴增根本不是TF的锅，是你把整个训练数据集一次性塞进内存了！换成分批加载的方式，比如用tf.data.Dataset的分批功能，或者Keras的Sequence生成器：

from keras.utils import Sequence
import numpy as np

class DataGenerator(Sequence):
    def __init__(self, data_paths, batch_size=32):
        self.data_paths = data_paths
        self.batch_size = batch_size

    def __len__(self):
        return int(np.ceil(len(self.data_paths) / self.batch_size))

    def __getitem__(self, idx):
        # 每次只加载一批数据，而不是全量
        batch_paths = self.data_paths[idx*self.batch_size:(idx+1)*self.batch_size]
        batch_data = [load_and_preprocess(path) for path in batch_paths]
        return np.array(batch_data)

这样内存里只会保留当前批次的数据，压力瞬间小很多。

3. 用cgroups替代ulimit，精准限制进程内存 #

ulimit对TF这类进程的内存限制效果真的有限，换成Linux的cgroups工具能精准控住内存上限：

• 先创建一个内存限制组：

  sudo cgcreate -g memory:/tf_limit
  

• 设置内存上限（比如设为6GB，单位是字节，6GB=6442450944）：

  sudo cgset -r memory.limit_in_bytes=6442450944 tf_limit
  

• 最后在这个限制组里启动训练脚本：

  sudo cgexec -g memory:tf_limit python3 your_train_script.py
  

这样就能强制TF最多用6GB内存，再也不会触发OOM了。

4. 检查模型本身是否太“臃肿” #

如果上面的方法都不管用，可能是你的模型结构本身太大：

• 先把批量大小（batch size）砍半，比如从64降到32甚至16
• 换掉超大的全连接层，或者用轻量化模型（比如MobileNet替代VGG）
• TF 2.x可以开启混合精度训练：tf.keras.mixed_precision.set_global_policy('mixed_float16')，能大幅降低内存占用

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