Keras训练微调BERT模型GPU新增全连接层触发OOM错误求助

你触发OOM的核心原因是：新增Dense层时传入的是完整sequence_output，其形状为[batch_size, max_len, bert_hidden_size]（你所用的BERT-large隐层维度为1024），新增Dense层会对序列的每个位置做映射，额外产生大量中间张量与梯度存储需求，刚好超出当前GPU的显存余量，因此只有加了这两层后才会触发OOM。

1. 特征输入优化（优先级最高） #

如果你的任务是文本分类/回归类非序列标注任务，不要传入完整sequence_output，改用BERT输出的pooled_output（即<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>特殊字符对应的聚合输出，形状为[batch_size, 1024]），可直接降低一个维度的显存占用，修改代码如下：

# 把原来的x = Dense(128, activation = "relu")(sequence_output)
# 替换为
x = Dense(128, activation = "relu")(pooled_output)

如果是序列标注/NER类必须用完整序列输出的任务，直接跳到后续方案。

2. 超参数调整 #

• 降低批次大小：将当前batch_size=16下调到8/4/2，批次过小的话可适当将学习率从2e-6调高到3e-6~5e-6，抵消小批次带来的训练波动。
• 降低序列最大长度：如果你的输入文本平均长度远小于256，可将max_len下调到128/64，序列长度减半显存占用也会对应减半。

3. 启用TensorFlow显存优化 #

3.1 开启显存动态分配 #

在代码最开头加入以下配置，避免TensorFlow启动时直接占满所有GPU显存：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

3.2 启用混合精度训练 #

开启混合精度后，TensorFlow会自动用float16存储大部分张量，显存占用可降低约50%，同时训练速度也会提升，修改编译代码即可（注意原代码中lr参数已被弃用，替换为learning_rate）：

from tensorflow.keras.mixed_precision import LossScaleOptimizer
opt = LossScaleOptimizer(tf.keras.optimizers.Adam(learning_rate = 2e-6))
model.compile(opt, loss = 'binary_crossentropy', metrics = ['accuracy'])

4. 模型结构精简 #

如果以上方案都无法满足需求，可将新增的Dense层维度从128下调到64/32，或直接删掉额外的Dense+Dropout层，直接在输入特征上接输出层。

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