TensorFlow调整CIFAR10图像尺寸时出现CPU OOM错误，求内存优化方案

这个问题我太熟悉了——一次性对整个CIFAR-10训练集执行resize操作，直接把CPU内存撑爆太正常了。毕竟50000张244×244×3的浮点张量算下来得占几十GB内存，普通机器根本顶不住。下面几个实用方法帮你降低内存占用：

1. 用TensorFlow Data Pipeline惰性处理（最推荐） #

tf.data.Dataset是TensorFlow官方推荐的数据流处理方式，它采用惰性加载+并行预处理的机制，不会一次性把所有图像都加载到内存里，每次只处理当前批次的图像，内存占用极低。

代码示例：

import tensorflow as tf

# 加载原始数据集
cifar10 = tf.keras.datasets.cifar10
(train_images, train_labels), (test_images, test_labels) = cifar10.load_data()

# 构建训练集数据流，在管道中实时resize
train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
# 并行执行resize，提高效率
train_dataset = train_dataset.map(
    lambda img, lbl: (tf.image.resize(img, (244, 244)), lbl),
    num_parallel_calls=tf.data.AUTOTUNE
)
# 分批次+预取，进一步优化训练效率
train_dataset = train_dataset.batch(32).prefetch(tf.data.AUTOTUNE)

# 测试集同理
test_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels))
test_dataset = test_dataset.map(
    lambda img, lbl: (tf.image.resize(img, (244, 244)), lbl),
    num_parallel_calls=tf.data.AUTOTUNE
)
test_dataset = test_dataset.batch(32).prefetch(tf.data.AUTOTUNE)

之后训练模型时直接传入train_dataset和test_dataset即可，全程不会出现内存过载问题。

2. 手动分批次处理 #

如果你确实需要把所有resize后的图像存入内存，可以分小批次逐步处理，每次只在内存中保留一个批次的图像，处理完再拼接结果：

代码示例：

import tensorflow as tf
import numpy as np

cifar10 = tf.keras.datasets.cifar10
(train_images, train_labels), (test_images, test_labels) = cifar10.load_data()

# 设置合理的批次大小（根据你的内存情况调整，比如1000）
batch_size = 1000
resized_train = []

# 分批次处理训练集
for i in range(0, len(train_images), batch_size):
    # 取出当前批次
    batch_imgs = train_images[i:i+batch_size]
    # 对批次图像执行resize
    resized_batch = tf.image.resize(batch_imgs, (244, 244)).numpy()
    resized_train.append(resized_batch)

# 拼接所有批次结果
train_images_resized = np.concatenate(resized_train, axis=0)

# 测试集同样处理
resized_test = []
for i in range(0, len(test_images), batch_size):
    batch_imgs = test_images[i:i+batch_size]
    resized_batch = tf.image.resize(batch_imgs, (244, 244)).numpy()
    resized_test.append(resized_batch)
test_images_resized = np.concatenate(resized_test, axis=0)

这种方法适合需要后续对resized图像做其他numpy操作的场景，但注意最终拼接后的结果还是会占用大量内存，如果内存依然不够，可以考虑用float16精度替代float32（在resize时指定dtype=tf.float16），能减少一半内存占用。

3. 自定义Keras数据生成器 #

如果你习惯用Keras的生成器接口，可以自定义一个Sequence类，每次生成批次时实时resize图像：

代码示例：

from tensorflow.keras.utils import Sequence
import tensorflow as tf
import numpy as np

class CIFAR10ResizeGenerator(Sequence):
    def __init__(self, images, labels, batch_size, target_size=(244,244)):
        self.images = images
        self.labels = labels
        self.batch_size = batch_size
        self.target_size = target_size
        self.indexes = np.arange(len(self.images))

    def __len__(self):
        # 返回总批次数量
        return len(self.images) // self.batch_size

    def __getitem__(self, idx):
        # 获取当前批次的索引
        batch_idx = self.indexes[idx*self.batch_size : (idx+1)*self.batch_size]
        # 取出批次图像和标签
        batch_imgs = self.images[batch_idx]
        batch_labels = self.labels[batch_idx]
        # 实时resize
        resized_imgs = tf.image.resize(batch_imgs, self.target_size).numpy()
        return resized_imgs, batch_labels

# 使用生成器
train_generator = CIFAR10ResizeGenerator(train_images, train_labels, batch_size=32)
test_generator = CIFAR10ResizeGenerator(test_images, test_labels, batch_size=32)

训练时直接把生成器传入model.fit()即可，内存占用和tf.data管道类似。

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