我来帮你排查下训练时CPU内存耗尽的问题，这在处理大型图像数据集时非常常见，大概率是数据加载和预处理环节的内存管理没做好，给你几个针对性的优化方案：

你提到训练前把数据集加载成4维NumPy数组，这可能是问题根源——如果数据集本身很大，一次性全量加载到CPU内存里，哪怕batch_size只有40，也会直接占满内存，甚至触发交换内存使用。另外，预处理过程中产生的临时变量未及时释放、数据格式冗余也会加剧内存压力。

1. 用迭代器/生成器分批加载数据（最关键） #

不要一次性把整个数据集读入内存，而是用按需加载的方式，每次只读取当前训练需要的batch数据。利用h5py的特性，它支持直接对磁盘上的数据集做切片操作，不用全量加载：

def train_data_generator(file_path, batch_size=40):
    with h5py.File(file_path, 'r') as _data:
        train_data = _data['train_data']  # 这里只是获取数据集引用，不加载到内存
        total_samples = train_data.shape[0]
        
        for start_idx in range(0, total_samples, batch_size):
            end_idx = min(start_idx + batch_size, total_samples)
            # 仅加载当前batch的数据到内存
            batch = train_data[start_idx:end_idx]
            # 在这里加入你的预处理逻辑（归一化、转置等）
            batch = batch.astype(np.float32) / 255.0  # 示例：转float32并归一化
            yield batch

训练时直接循环这个生成器即可，每次迭代只会把当前batch的数据留在内存里。

2. 优化预处理流程，避免内存泄漏 #

预处理环节很容易产生大量临时数组，一定要及时清理：

• 尽量用原地操作减少临时变量，比如batch /= 255.0代替batch = batch / 255.0
• 用del手动删除不再使用的变量，再调用gc.collect()触发垃圾回收：

# 预处理后清理临时变量
del temp_array
import gc
gc.collect()

3. 转换为更高效的数据集格式 #

.mat格式虽然通用，但对于大型数据集来说不够高效，建议转成支持内存映射的格式：

• NumPy内存映射文件：把数据集转成.npy后，用mmap_mode='r'加载，数据会留在磁盘，按需读取：

# 先转换格式（仅需执行一次）
from scipy.io import loadmat
raw_data = loadmat('traindata.mat')['train_data']
np.save('traindata.npy', raw_data)

# 训练时加载
train_data = np.load('traindata.npy', mmap_mode='r')
# 取batch的方式和h5py一致：train_data[start:end]

• 也可以考虑转成TFRecord（TensorFlow）或LMDB，这类格式专门针对大样本训练做了优化，支持随机访问和分批读取。

4. 压缩数据体积 #

• 降低数据精度：如果当前用的是float64，转成float32能直接减少一半内存占用；图像数据甚至可以用uint8存储，预处理时再转成float32
• 缩小图像尺寸：如果图像分辨率过高（比如2048x2048），可以提前resize到模型需要的大小（比如512x512），内存占用会大幅降低，对模型性能影响通常很小

5. 检查训练框架的内存配置 #

如果你用TensorFlow/PyTorch这类框架，确保没有开启不必要的内存预分配：

• TensorFlow：可以设置tf.config.experimental.set_memory_growth(gpu, True)，避免GPU预占过多内存，但更关键的是CPU端的数据加载逻辑
• PyTorch：用DataLoader时设置num_workers不要过大，否则多个子进程会占用额外CPU内存

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