TensorFlow CNN训练遇OutOfMemory错误，求数据缩减解决办法

Hey there, let's tackle this OutOfMemory issue you're facing with your TensorFlow CNN model when adapting it to your custom 244×244 data. Given your Ubuntu 14.04 setup with a 4-core CPU and 16GB RAM, here are several practical strategies to reduce memory usage and get your training up and running:

1. 进一步减小Batch Size #

Your current batch size is set to 5, which might still be too large for 244×244 inputs. Try scaling it down even more—start with 2 or 1 to see if that resolves the OOM error.

Modify your train_input_fn like this:

train_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={"x": np.array(training_set.data)},
    y=np.array(training_set.target),
    num_epochs=None,
    batch_size=2,  # 尝试更小的batch size
    shuffle=True
)

2. 缩小输入图像尺寸 #

244×244 is a relatively large input size. If your task allows, resize your images to a smaller dimension (like 128×128 or 64×64) before feeding them into the model. This drastically reduces the memory footprint of each input sample.

You can preprocess your data using libraries like PIL or OpenCV:

# 示例：用PIL调整图像尺寸
from PIL import Image
import numpy as np

# 假设training_set.data是形状为(n_samples, 244, 244, channels)的数组
resized_data = []
for img in training_set.data:
    pil_img = Image.fromarray(img)
    resized_img = pil_img.resize((128, 128))  # 调整为128×128
    resized_data.append(np.array(resized_img))
resized_data = np.array(resized_data)

# 然后在input_fn中使用resized_data
train_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={"x": resized_data},
    y=np.array(training_set.target),
    num_epochs=None,
    batch_size=5,
    shuffle=True
)

3. 使用更轻量化的CNN结构 #

If your current model has too many layers or filters, trimming it down can save a lot of memory. For example:

• Reduce the number of filters in convolutional layers (e.g., from 64 to 32)
• Remove unnecessary fully connected layers or reduce their neuron count

Modify your cnn_model_fn accordingly—here's a simplified example:

def cnn_model_fn(features, labels, mode):
    # 输入层
    input_layer = tf.reshape(features["x"], [-1, 244, 244, 3])  # 假设是3通道图像

    # 卷积层1：减少滤波器数量
    conv1 = tf.layers.conv2d(
        inputs=input_layer,
        filters=32,  # 从64改为32
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu
    )
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)

    # 卷积层2：同样减少滤波器
    conv2 = tf.layers.conv2d(
        inputs=pool1,
        filters=64,  # 从128改为64
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu
    )
    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)

    # 全连接层：减少神经元数量
    pool2_flat = tf.reshape(pool2, [-1, 61*61*64])  # 244/2/2=61
    dense = tf.layers.dense(inputs=pool2_flat, units=256, activation=tf.nn.relu)  # 从512改为256
    dropout = tf.layers.dropout(
        inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)

    # 输出层
    logits = tf.layers.dense(inputs=dropout, units=10)  # 假设是10分类任务

    # 后续的模型逻辑（预测、损失、训练操作）保持不变
    # ...

4. 优化CPU内存配置 #

Even on CPU, you can configure TensorFlow to use memory more efficiently. Add a RunConfig when creating your Estimator to set CPU thread limits and avoid excessive memory consumption:

run_config = tf.estimator.RunConfig(
    session_config=tf.ConfigProto(
        intra_op_parallelism_threads=2,
        inter_op_parallelism_threads=2,
        allow_soft_placement=True
    )
)

mnist_classifier = tf.estimator.Estimator(
    model_fn=cnn_model_fn,
    model_dir="path/to/model",
    config=run_config
)

This prevents TensorFlow from spawning too many threads, which can eat up memory on your 4-core CPU.

5. 使用数据生成器加载数据（避免一次性加载全部数据） #

Right now, you're loading your entire dataset into memory with np.array(training_set.data). If your dataset is large, this alone can cause OOM. Instead, use a generator to load batches of data on-the-fly:

def data_generator():
    while True:
        # 每次生成一个batch的数据
        for i in range(0, len(training_set.data), 5):
            batch_x = training_set.data[i:i+5]
            # 可选：在这里调整图像尺寸
            batch_x = np.array([np.array(Image.fromarray(img).resize((128,128))) for img in batch_x])
            batch_y = training_set.target[i:i+5]
            yield {"x": batch_x}, batch_y

train_input_fn = tf.estimator.inputs.numpy_input_fn(
    x=None,
    y=None,
    input_fn=data_generator,
    num_epochs=None,
    shuffle=True
)

This way, only one batch of data is loaded into memory at a time, drastically reducing overall memory usage.

你的原始代码参考：

# Create the Estimator
mnist_classifier = tf.estimator.Estimator(
    model_fn=cnn_model_fn,
    model_dir="path/to/model")
# Load the data
train_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={"x": np.array(training_set.data)},
    y=np.array(training_set.target),
    num_epochs=None,
    batch_size=5,
    shuffle=True)
# Train the model
mnist_classifier.train(
    input_fn=train_input_fn,
    steps=100,
    hooks=[logging_hook])

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