Google Colab运行RoBERTa推特情感提取时出现ResourceExhaustedError如何解决

这个报错是GPU显存耗尽导致的显存分配失败，当前配置的显存需求超过了Colab平台GPU的可用上限：

• roberta-large模型本身有355M参数，全量参与训练时前向、反向传播的显存占用基数很高
• 报错张量的第一个维度为32，对应训练时的batch size是32，搭配128长度的输入序列，注意力层的显存占用会成倍放大
• Colab免费版GPU的可用显存通常为16G，无法支撑当前配置的显存需求

• 降低训练batch size
  把当前训练用的batch size从32降到8，如果仍然OOM就进一步降到4，直到显存占用符合上限。对应调整tf.data管道的batch()参数、或model.fit()接口的batch_size参数即可
• 替换为RoBERTa-Base权重
  将TFRobertaModel.from_pretrained('roberta-large')中的权重名改为roberta-base，模型参数量直接降到125M，显存占用相比roberta-large降低60%以上，对短文本情感抽取任务的精度影响极小
• 开启梯度检查点
  加载预训练模型时增加梯度检查点参数，把部分前向传播的中间存储转成反向传播时重算，可降低30%-40%的显存占用，训练速度会有一定损耗但可正常运行：

  bert_model = TFRobertaModel.from_pretrained('roberta-large', gradient_checkpointing=True)
  

• 缩短最大输入序列长度
  推特文本普遍长度较短，可将MAX_LEN从当前128降到64，注意力机制的显存占用和序列长度的平方正相关，序列长度减半可减少75%的注意力层显存占用
• 开启混合精度训练
  在脚本最开头（加载所有模型之前）配置全局混合精度策略，将大部分计算从32位浮点数转为16位，直接降低一半的计算显存占用，Colab GPU完全兼容该特性，几乎无精度损失：

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

• 冻结预训练模型大部分层
  固定预训练模型的底层参数，仅微调最后2-3层，减少反向传播时需要存储的梯度显存：

  bert_model.trainable = False
  # 解冻最后2层进行微调
  for layer in bert_model.layers[-2:]:
      layer.trainable = True
  

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