16GB显存P5000训练Keras模型显存不足问题排查

环境与任务概况 #

• 平台：Paperspace，显卡为16GB显存的P5000
• 框架：Keras（基于TensorFlow）
• 模型：约800万可训练参数的图像分割模型
• 数据集：Kaggle Carvana图像分割数据集，含约5000张图像，原图尺寸1918x1280，训练时缩至512x512
• 数据加载：使用TensorFlow Dataset API，预处理代码如下：

import os
import math
import random
from typing import Tuple
import tensorflow as tf

def randomize_image(image, seed: int, apply_color_changes: bool):
    image = tf.image.random_flip_left_right(image, seed=seed)
    image = tf.image.random_flip_up_down(image, seed=seed)
    image = tf.keras.layers.RandomRotation(factor = 0.2,seed=seed)(image)
    if apply_color_changes:
        image = tf.image.random_contrast(image, lower=.4, upper=.6,seed=seed)
        image = tf.image.random_brightness(image, max_delta=.2,seed=seed)
        image = tf.image.random_saturation(image, lower=.4, upper=.6,seed=seed)
        image = tf.image.random_hue(image, max_delta=.2,seed=seed)
    return image

def randomize_images(image, mask):
    seed = random.randint(0, 100000000)
    image = randomize_image(image, seed, True)
    mask = randomize_image(mask, seed, False)
    return image, mask

def resize(image, mask, input_shape: Tuple[int, int]):
    image = tf.image.resize(image,input_shape)
    mask = tf.image.resize(mask,input_shape)
    return image / 255, mask / 255

def load_images(image, mask):
    image = tf.io.read_file(image)
    image = tf.io.decode_jpeg(image)
    mask = tf.io.read_file(mask)
    mask = tf.io.decode_image(mask, channels=0, expand_animations = False)
    mask = tf.image.rgb_to_grayscale(mask)
    return image, mask


def get_dataset(image_dir, mask_dir, input_shape: Tuple[int, int], randomize_images=True, batch_size: int = 16, val_split: float = .05) -> tf.data.Dataset:
    images = os.listdir(image_dir)
    masks = [image.replace(".jpg", "_mask.gif") for image in images]
    images = [os.path.join(image_dir,image) for image in images]
    masks = [os.path.join(mask_dir,mask) for mask in masks]
    data = list(zip(images,masks))
    random.shuffle(data)
    
    train_size = math.floor(len(images) * (1-val_split))

    train_data = data[:train_size]
    val_data = data[train_size:]
    train_images, train_masks = zip(*train_data)
    val_images, val_masks = zip(*val_data)

    train_dataset = tf.data.Dataset.from_tensor_slices((list(train_images), list(train_masks))).map(map_func=load_images, num_parallel_calls=tf.data.AUTOTUNE)
    val_dataset = tf.data.Dataset.from_tensor_slices((list(val_images), list(val_masks))).map(map_func=load_images, num_parallel_calls=tf.data.AUTOTUNE)
    if randomize_images:
        train_dataset = train_dataset.map(map_func=randomize_images, num_parallel_calls=tf.data.AUTOTUNE)
    
    train_dataset = train_dataset.map(map_func=lambda image, mask: resize(image, mask, input_shape), num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size)
    val_dataset = val_dataset.map(map_func=lambda image, mask: resize(image, mask, input_shape), num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size)


    return train_dataset, val_dataset

报错信息 #

Allocator (GPU_0_bfc) ran out of memory trying to allocate 3.39GiB with freed_by_count=0

显存不足的核心原因 #

• 批量数据+训练过程的总显存远超原始数据计算值：
  单张512x512 RGB图像（float32）仅占约3.1MiB，对应掩码占1.0MiB，32个样本的原始数据仅约133MiB，但训练时还有大量额外开销：

  • 模型中间特征图：图像分割模型（如U-Net结构）会生成多尺度特征图，每层特征图的通道数可能达到64、128甚至256，这部分显存占用是原始数据的数倍甚至数十倍。
  • 反向传播梯度：除了800万参数的梯度（约30MiB），中间层激活值的梯度才是显存占用的大头，尤其是大尺寸特征图的梯度张量。
  • 优化器状态：比如Adam优化器会保存每个参数的一阶矩和二阶矩，相当于额外占用2倍参数显存（约61MiB）。

• 数据预处理的潜在显存浪费：
  randomize_images中混用Python原生random.randint和TensorFlow操作，可能导致Eager模式下临时张量未及时释放；同时tf.keras.layers.RandomRotation作为层在tf.data.map中使用，可能在计算图中残留不必要的张量节点，增加显存占用。

• 平台与TensorFlow显存策略限制：
  Paperspace平台的后台服务、监控进程会占用部分显存，导致实际可用显存少于标称的16GB；TensorFlow默认预分配大部分显存，加上显存碎片问题，会进一步压缩可用空间。当batch size超过32时，额外的张量需求刚好突破剩余显存阈值。

• 验证集的叠加占用：
  训练时验证集的batch会同步加载，若验证集batch size与训练集一致，相当于同时持有两个完整batch的数据，进一步加剧显存压力。

额外排查方向 #

• 检查是否开启了TensorFlow调试工具（如tf.debugging.experimental.enable_dump_debug_info），这类工具会保存大量中间张量，占用额外显存。
• 确认模型是否存在冗余输出或层，比如不必要的辅助分支，会增加张量占用。

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