TensorFlow模型训练内存溢出调试求助：单轮epoch后崩溃

环境与基本情况 #

• 模型参数大小：230MB（总参数60130177）
• 数据集规模：300MB
• 任务：二分类CNN训练
• 系统配置：16GB内存 + RTX 4070Ti显卡
• 核心问题：单轮Epoch后崩溃，即使将batch size降至1也无法解决，目标是支持batch size=64

错误提示 #

102/102 [==============================] - ETA: 0s - loss: 1.0180 - accuracy: 0.6293 - precision: 0.0000e+00 - recall: 0.0000e+002024-04-12 08:04:29.166412: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 12582912000 exceeds 10% of free system memory.
Killed

模型代码 #

def create_dual_stream_cnn_model(input_shape):
    # 定义输入
    input = Input(shape=input_shape)
    
    # 分支1
    x = Conv1D(64, 3, activation='relu', padding='same')(input)
    x = Conv1D(64, 3, activation='relu', padding='same')(x)
    x = MaxPooling1D(3, strides=3)(x)

    x = Conv1D(128, 3, activation='relu', padding='same')(x)
    x = Conv1D(128, 3, activation='relu', padding='same')(x)
    x = MaxPooling1D(3, strides=3)(x)

    x = Conv1D(256, 3, activation='relu', padding='same')(x)
    x = Conv1D(256, 3, activation='relu', padding='same')(x)
    x = MaxPooling1D(2, strides=2)(x)

    x = Conv1D(512, 3, activation='relu', padding='same')(x)
    x = Conv1D(512, 3, activation='relu', padding='same')(x)
    x = MaxPooling1D(2, strides=2)(x)

    x = Conv1D(512, 3, activation='relu', padding='same')(x)
    x = Conv1D(512, 3, activation='relu', padding='same')(x)
    x = MaxPooling1D(2, strides=2)(x)
    
    # 分支2
    y = Conv1D(64, 7, activation='relu', padding='same')(input)
    y = Conv1D(64, 7, activation='relu', padding='same')(y)
    y = MaxPooling1D(3, strides=3)(y)
    
    y = Conv1D(128, 7, activation='relu', padding='same')(y)
    y = Conv1D(128, 7, activation='relu', padding='same')(y)
    y = MaxPooling1D(3, strides=3)(y)

    y = Conv1D(256, 3, activation='relu', padding='same')(y)
    y = Conv1D(256, 3, activation='relu', padding='same')(y)
    y = MaxPooling1D(2, strides=2)(y)

    y = Conv1D(512, 3, activation='relu', padding='same')(y)
    y = Conv1D(512, 3, activation='relu', padding='same')(y)
    y = MaxPooling1D(2, strides=2)(y)
    
    y = Conv1D(512, 3, activation='relu', padding='same')(y)
    y = Conv1D(512, 3, activation='relu', padding='same')(y)
    y = MaxPooling1D(2, strides=2)(y)
    
    # 拼接两个分支
    concatenated = concatenate([x, y])

    z = Flatten()(concatenated)
    z = Dense(1024, activation='relu', kernel_regularizer=L2(0.0001))(z)
    z = Dense(1024, activation='relu', kernel_regularizer=L2(0.0001))(z)
    z = Dense(256, activation='relu', kernel_regularizer=L2(0.0001))(z)
    z = Dense(1, activation='sigmoid')(z)
    
    model = Model(inputs=input, outputs=z)
    optimizer = SGD()
    metrics = ['accuracy', 'Precision', 'Recall']
    model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=metrics)
    print(model.summary())
    return model

训练循环代码 #

BATCH_SIZE = 64
EPOCHS = 400
K_FOLDS = 10

X = np.array(cropped_records)
y = np.array(dup_labels)
y = y[:,0].astype(int)
X = np.expand_dims(X, -1)

kf = KFold(n_splits=K_FOLDS, shuffle=True)

test_scores = []
fold_id = 0
train_time = datetime.now().strftime("%Y%m%d_%H%M%S")

for train_index, test_index in kf.split(X):

    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]

    X_train, X_valid, y_train, y_valid = train_test_split(X_train, y_train, test_size=0.1)
    train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
    validation_dataset = tf.data.Dataset.from_tensor_slices((X_valid, y_valid))
    test_dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test))

    train_dataset = train_dataset.shuffle(buffer_size=100).batch(BATCH_SIZE).prefetch(buffer_size=BATCH_SIZE*3)
    validation_dataset = validation_dataset.batch(BATCH_SIZE).prefetch(buffer_size=BATCH_SIZE*3)
    test_dataset = test_dataset.batch(BATCH_SIZE).prefetch(buffer_size=BATCH_SIZE*3)

    logs_dir = 'logs/' + train_time + f'/{fold_id}'
    if not os.path.exists(logs_dir):
        os.makedirs(logs_dir)

    model = create_dual_stream_cnn_model((X_train.shape[1], 1))
    print_gpu_availability()

    tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logs_dir, histogram_freq=1)
    lr_scheduler = tf.keras.callbacks.LearningRateScheduler(exponential_decay_fn)
    model.fit(train_dataset, 
            epochs=EPOCHS, verbose=1,
            callbacks=[lr_scheduler, tensorboard_callback])

    test_loss, test_accuracy, test_precision, test_recall = model.evaluate(test_dataset)
    y_scores = model.predict(X_test, verbose=0)
    y_scores = y_scores.flatten()
    test_fpr, test_tpr, _ = roc_curve(y_test, y_scores)
    test_auc = roc_auc_score(y_test, y_scores)

    test_scores.append({'loss':test_loss, 
                        'acc': test_accuracy, 
                        'prec':test_precision, 
                        'rec':test_recall, 
                        'auc':test_auc, 
                        'fpr':test_fpr, 
                        'tpr':test_tpr})

    fold_id += 1

1. 内存溢出根源 #

错误提示显示尝试分配12.5GB内存到CPU，核心原因是：

• TensorBoard的histogram_freq=1会存储每层权重和激活值的直方图，占用大量内存
• 模型Flatten后的全连接层张量维度过高，加上反向传播的梯度、优化器状态等，导致内存占用远超预期
• K折循环中未清理旧模型，内存持续累积

2. 针对性优化措施 #

（1）关闭TensorBoard直方图记录

将tensorboard_callback的histogram_freq设为0，避免存储大量直方图数据：

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logs_dir, histogram_freq=0)

（2）添加内存清理步骤

在每折训练结束后清理模型占用的内存，避免累积：

# 每折末尾添加
del model
tf.keras.backend.clear_session()
fold_id += 1

（3）优化模型结构，降低张量维度

用GlobalAveragePooling1D替代Flatten，大幅减少全连接层输入维度：

# 替换原Flatten部分
x = GlobalAveragePooling1D()(x)
y = GlobalAveragePooling1D()(y)
concatenated = concatenate([x, y])
# 可同步缩小全连接层规模，减少参数和内存占用
z = Dense(512, activation='relu', kernel_regularizer=L2(0.0001))(concatenated)
z = Dense(256, activation='relu', kernel_regularizer=L2(0.0001))(z)
z = Dense(1, activation='sigmoid')(z)

（4）限制CPU内存增长

在代码开头添加TensorFlow内存配置，避免无限制占用CPU内存：

import tensorflow as tf
# 开启GPU内存动态增长
physical_devices = tf.config.list_physical_devices('GPU')
if physical_devices:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)
# 限制CPU内存占用为8GB
tf.config.set_logical_device_configuration(
    tf.config.list_physical_devices('CPU')[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=8192)]
)

（5）检查输入序列长度

若输入序列过长（如超过10000），可截断或下采样，进一步降低Flatten/池化后的张量维度。

3. 验证顺序 #

先尝试关闭直方图记录+添加内存清理，这两个操作无需修改模型结构，可快速验证是否解决问题；若仍有内存压力，再调整模型结构替换Flatten为全局池化。

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