如何解决TensorFlow训练中的ResourceExhaustedError资源耗尽错误

一、训练参数调整 #

• 降低Batch Size：当前batch size=16，GTX1650显存容量有限（多数版本为4G），建议先将batch size降至8甚至4，这是最直接的缓解方式。可根据训练时的显存占用情况逐步测试，找到能稳定运行的最大值。
• 启用梯度累积：若不想降低batch size，可通过梯度累积模拟大batch训练效果。每次计算小batch的梯度但不更新权重，累积指定次数后再一次性更新权重。示例代码：

accumulation_steps = 2  # 累积2次等价于batch size=32
overall_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}")
    total_loss = 0.0
    steps = 0
    accumulated_grads = None
    
    for x_batch, y_batch in train_dataset:
        with tf.GradientTape() as tape:
            predictions = overall_model(x_batch, training=True)
            loss = overall_model.compiled_loss(y_batch, predictions)
        
        grads = tape.gradient(loss, overall_model.trainable_variables)
        # 累积梯度
        if accumulated_grads is None:
            accumulated_grads = grads
        else:
            for i in range(len(accumulated_grads)):
                accumulated_grads[i] += grads[i]
        
        steps += 1
        if steps % accumulation_steps == 0:
            # 应用累积梯度更新权重
            overall_model.optimizer.apply_gradients(zip(accumulated_grads, overall_model.trainable_variables))
            accumulated_grads = None
            total_loss += loss.numpy()
            print(f"Step {steps}, Loss: {loss.numpy():.4f}")
    
    avg_loss = total_loss / (steps // accumulation_steps) if steps // accumulation_steps > 0 else 0
    print(f"Epoch {epoch+1} Average Loss: {avg_loss:.4f}")

二、数据加载优化 #

• 将预处理放入Dataset管道：resize_and_rescale和data_augmentation若在内存中批量处理，会占用大量显存。建议通过tf.data.Dataset的map方法并行预处理，并启用预取，避免一次性加载所有预处理后的数据：

train_dataset = train_dataset.map(lambda x, y: (resize_and_rescale(x), y), num_parallel_calls=tf.data.AUTOTUNE)
train_dataset = train_dataset.map(lambda x, y: (data_augmentation(x, training=True), y), num_parallel_calls=tf.data.AUTOTUNE)
train_dataset = train_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)

• 缩小输入图片尺寸：若input_shape过大（如(224,224,3)及以上），会大幅增加显存占用。可尝试将尺寸缩小至(128,128,3)，这对多数分类任务的精度影响有限，但能显著降低显存压力。

三、模型结构优化 #

• 减少卷积通道数：当前模型的卷积层通道数为32→64→64，可适当缩减为16→32→32，减少参数总量和中间特征图的显存占用：

overall_model = models.Sequential([
    resize_and_rescale, data_augmentation,
    layers.Conv2D(16, (3,3), activation="relu", input_shape=input_shape),
    layers.MaxPooling2D((2,2)),
    layers.Conv2D(32, kernel_size=(3,3), activation="relu"),
    layers.MaxPooling2D((2,2)),
    layers.Conv2D(32, kernel_size=(3,3), activation="relu"),
    layers.MaxPooling2D((2,2)),
    layers.Flatten(),
    layers.Dense(32, activation='relu'),
    layers.Dense(len(class_names), activation='softmax')
])

• 添加Dropout层：在Flatten之后或全连接层之间添加Dropout，既能防止过拟合，又能减少训练时的显存占用（随机失活部分神经元）：

layers.Flatten(),
layers.Dropout(0.2),
layers.Dense(64, activation='relu'),
layers.Dropout(0.2),
layers.Dense(len(class_names), activation='softmax')

• 启用混合精度训练：TensorFlow的混合精度可自动将部分运算转为FP16格式，大幅降低显存占用。只需添加以下代码：

from tensorflow.keras.mixed_precision import set_global_policy
set_global_policy('mixed_float16')

注意最后一层全连接层需保持float32输出，避免精度损失：

layers.Dense(len(class_names), activation='softmax', dtype='float32')

四、其他显存优化技巧 #

• 训练前清理显存：在初始化模型前添加tf.keras.backend.clear_session()，释放之前模型残留的显存：

import tensorflow as tf
tf.keras.backend.clear_session()

• 降低TensorBoard写入频率：若开启了TensorBoard实时监控，高频率写入会占用额外显存，可改为每几个epoch写入一次，或训练结束后再生成日志。

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