使用NumPy concatenate合并数据时遇ArrayMemoryError的技术问询

问题根源 #

当前报错的核心是float64数据类型内存占用过高：单张224×224×3的float64图片占1.17MB，9000张总计约10GB，加上系统本身占用和加载过程中的临时内存，14GB空闲内存不足以支撑。同时原代码先把所有图片加载到列表再合并，会产生两倍内存峰值（列表中所有数组+合并后的大数组）。

解决方法 #

1. 降低数据类型（最快见效）

深度学习中float32完全满足精度需求，将数据类型从float64转为float32可直接减半内存占用，若模型支持（如TensorFlow混合精度训练），还可进一步用float16：

# 加载时直接转换数据类型
cats_train_data = [np.load("E:\Unity\!!neuro\datasets\catsAndDogs100\finishedCats1\\" + filename).astype(np.float32) 
                   for filename in os.listdir("E:\Unity\!!neuro\datasets\catsAndDogs100\finishedCats1\\")]
dogs_train_data = [np.load("E:\Unity\!!neuro\datasets\catsAndDogs100\finishedDogs1\\" + filename).astype(np.float32) 
                   for filename in os.listdir("E:\Unity\!!neuro\datasets\catsAndDogs100\finishedDogs1\\")]

# 后续合并代码不变
train_data = np.concatenate((cats_train_data, dogs_train_data), axis=0)
train_labels = np.concatenate((cats_train_labels, dogs_train_labels), axis=0)

2. 用TensorFlow Dataset分批加载（彻底解决内存问题）

无需一次性加载所有数据到内存，通过tf.data.Dataset实现按需分批加载，训练时仅将当前批次数据读入内存：

import tensorflow as tf
import os
import numpy as np

def load_sample(file_path, label):
    # 加载.npy文件并转换为Tensor
    img_np = np.load(file_path.numpy().decode('utf-8')).astype(np.float32)
    return tf.convert_to_tensor(img_np), tf.convert_to_tensor(label, dtype=tf.int32)

# 生成所有文件路径与对应标签
cat_dir = "E:\Unity\!!neuro\datasets\catsAndDogs100\finishedCats1\\"
dog_dir = "E:\Unity\!!neuro\datasets\catsAndDogs100\finishedDogs1\\"

cat_files = [os.path.join(cat_dir, f) for f in os.listdir(cat_dir)]
dog_files = [os.path.join(dog_dir, f) for f in os.listdir(dog_dir)]

all_files = cat_files + dog_files
all_labels = [0]*len(cat_files) + [1]*len(dog_files)

# 创建数据集并处理
dataset = tf.data.Dataset.from_tensor_slices((all_files, all_labels))
# 用py_function包装自定义加载逻辑
dataset = dataset.map(lambda x, y: tf.py_function(load_sample, [x, y], [tf.float32, tf.int32]),
                      num_parallel_calls=tf.data.AUTOTUNE)
# 设置批次大小、预取，提升训练效率
dataset = dataset.batch(32).prefetch(tf.data.AUTOTUNE)

# 训练时直接传入dataset
# model.fit(dataset, epochs=10, ...)

3. 预分配数组加载（减少内存峰值）

避免先将所有图片存入列表再合并，直接预分配最终大小的数组，逐个加载写入，降低内存峰值：

import numpy as np
import os

cat_dir = "E:\Unity\!!neuro\datasets\catsAndDogs100\finishedCats1\\"
dog_dir = "E:\Unity\!!neuro\datasets\catsAndDogs100\finishedDogs1\\"

cat_files = os.listdir(cat_dir)
dog_files = os.listdir(dog_dir)

total_samples = len(cat_files) + len(dog_files)
# 预分配float32类型的数组
train_data = np.empty((total_samples, 224, 224, 3), dtype=np.float32)
train_labels = np.empty(total_samples, dtype=np.int32)

# 加载猫的数据
for idx, filename in enumerate(cat_files):
    train_data[idx] = np.load(os.path.join(cat_dir, filename)).astype(np.float32)
    train_labels[idx] = 0

# 加载狗的数据
offset = len(cat_files)
for idx, filename in enumerate(dog_files):
    train_data[offset + idx] = np.load(os.path.join(dog_dir, filename)).astype(np.float32)
    train_labels[offset + idx] = 1

4. 缩小图片尺寸（进一步降低内存）

若上述方法仍不够，可将图片尺寸从224×224缩小到112×112，内存占用直接变为原来的1/4，猫狗分类任务中该尺寸足以保证精度：

# 加载时同时缩小尺寸（以步长2采样实现快速缩小）
cats_train_data = [np.load(os.path.join(cat_dir, filename)).astype(np.float32)[::2, ::2, :] 
                   for filename in cat_files]

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