删除变量后仍系统内存不足？大数据集分批训练求助

问题 #

我有一个无法全部载入内存的大型数据集，计划采用分批训练：每次载入一半数据集，训练2个epoch后删除该数据，再载入另一半重复流程。但即便每2个epoch后删除数据，还是会因系统内存不足崩溃。单独训练每一半都正常，循环执行就崩溃（不是GPU内存不足，是系统内存耗尽）。循环代码如下：

for i in range(20):
  if i % 2 == 0:
    X_train = np.load('/content/drive/My Drive/Kaggle ISLR/X_train_batch1.npy')
    y_train = np.load('/content/drive/My Drive/Kaggle ISLR/y_train_batch1.npy')
    X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42)
    train_gen = MyDataGenerator(X_train, y_train, batch_size, first_dim)
    val_gen = MyDataGenerator(X_val, y_val, batch_size, first_dim)
    del X_train
    del y_train
    model.fit(
    train_gen,
    validation_data=val_gen,
    epochs=2,
    callbacks=[checkpoint_callback]
    )
    del train_gen
    del val_gen
    
  else:
    X_train = np.load('/content/drive/My Drive/Kaggle ISLR/X_train_batch2.npy')
    y_train = np.load('/content/drive/My Drive/Kaggle ISLR/y_train_batch2.npy')
    X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42)
    train_gen = MyDataGenerator(X_train, y_train, batch_size, first_dim)
    val_gen = MyDataGenerator(X_val, y_val, batch_size, first_dim)
    del X_train
    del y_train
    model.fit(
    train_gen,
    validation_data=val_gen,
    epochs=2,
    callbacks=[checkpoint_callback]
    )
    del train_gen
    del val_gen

请问有没有更优方案避免内存不足？

优化方案 #

1. 手动触发垃圾回收 #

del只是删除变量引用，Python垃圾回收器不一定会立刻释放内存。在每次删除变量后强制触发回收，确保内存及时释放：

import gc

# 在每个分支的model.fit()之后添加
del X_train, y_train, train_gen, val_gen
gc.collect()

2. 用内存映射文件加载数据 #

不要用np.load()一次性加载整个npy文件，改用np.memmap将文件映射到内存，按需读取数据，避免一次性占满内存：

# 替换原有的np.load，注意指定和原文件一致的dtype
X_train = np.memmap('/content/drive/My Drive/Kaggle ISLR/X_train_batch1.npy', dtype='float32', mode='r')
y_train = np.memmap('/content/drive/My Drive/Kaggle ISLR/y_train_batch1.npy', dtype='int32', mode='r')

3. 修改数据生成器，直接从磁盘取数 #

让MyDataGenerator不再依赖已加载的内存数组，而是在取batch时才从磁盘读取对应数据，彻底避免内存堆积：

class MyDataGenerator(keras.utils.Sequence):
    def __init__(self, x_path, y_path, batch_size, first_dim, dtype_x='float32', dtype_y='int32'):
        self.x_path = x_path
        self.y_path = y_path
        self.batch_size = batch_size
        self.first_dim = first_dim
        # 仅读取数据形状，不加载全部数据
        with np.load(x_path, mmap_mode='r') as x:
            self.n_samples = x.shape[0]
        self.dtype_x = dtype_x
        self.dtype_y = dtype_y

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

    def __getitem__(self, idx):
        start = idx * self.batch_size
        end = min((idx+1)*self.batch_size, self.n_samples)
        # 每次仅读取当前batch的数据
        X = np.memmap(self.x_path, dtype=self.dtype_x, mode='r', shape=(self.n_samples, self.first_dim))[start:end]
        y = np.memmap(self.y_path, dtype=self.dtype_y, mode='r', shape=(self.n_samples,))[start:end]
        return X, y

4. 提前拆分验证集并保存 #

每次加载整个批次后执行train_test_split会额外占用内存，提前拆分好训练集和验证集并单独保存，训练时直接加载拆分后的文件：

# 提前执行一次，之后训练无需再拆分
# 处理batch1
X_batch1 = np.load('/content/drive/My Drive/Kaggle ISLR/X_train_batch1.npy')
y_batch1 = np.load('/content/drive/My Drive/Kaggle ISLR/y_train_batch1.npy')
X_tr1, X_val1, y_tr1, y_val1 = train_test_split(X_batch1, y_batch1, test_size=0.2, random_state=42)
np.save('/content/drive/My Drive/Kaggle ISLR/X_tr1.npy', X_tr1)
np.save('/content/drive/My Drive/Kaggle ISLR/y_tr1.npy', y_tr1)
np.save('/content/drive/My Drive/Kaggle ISLR/X_val1.npy', X_val1)
np.save('/content/drive/My Drive/Kaggle ISLR/y_val1.npy', y_val1)

# 处理batch2同理

5. 监控内存泄漏 #

用psutil库实时查看内存使用，定位哪一步内存未正常释放：

import psutil

def print_memory_usage():
    process = psutil.Process()
    print(f"当前内存使用: {process.memory_info().rss / 1024**2:.2f} MB")

# 在循环的关键节点（加载数据前、训练后、垃圾回收后）调用
print_memory_usage()

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