如何移除GPU显存中加载的Transformers模型并释放显存

关于torch.cuda.empty_cache()的适用性 #

torch.cuda.empty_cache()是PyTorch框架专属的GPU显存释放接口，你当前代码加载的是TensorFlow版本的CTRL模型（类名前缀TF为TensorFlow实现标识），调用这个接口无法释放你当前占用的GPU显存。

TensorFlow框架下释放GPU显存的正确步骤 #

按顺序执行以下操作即可完全释放相关显存：

• 第一步：删除所有关联的对象变量，切断引用关系

del gen_nlp, model_ctrl, tokenizer_ctrl

• 第二步：调用TensorFlow内置接口清空会话显存

import tensorflow as tf
# 清空后端会话缓存
tf.keras.backend.clear_session()
# 如果你要清理指定GPU（你之前用的是device=1即编号为1的GPU）的统计信息与残留显存，可以额外执行
tf.config.experimental.reset_memory_stats('GPU:1')

补充：PyTorch版本模型的释放方法 #

如果你后续切换为PyTorch实现的CTRL模型，可以按以下步骤操作：

• 第一步：删除所有关联变量，若模型/张量在GPU上，也可以手动移动到CPU后再删除

del gen_nlp, model_ctrl, tokenizer_ctrl

• 第二步：调用PyTorch显存清理接口释放未被占用的缓存显存

import torch
# 注意：该接口只会释放没有任何变量引用的显存缓存，如果还有变量持有GPU显存引用，调用后不会释放对应显存
torch.cuda.empty_cache()

注意事项 #

• 如果是在Jupyter Notebook/Colab环境中操作，需要确保没有其他单元格的输出、变量持有上述对象的引用，必要时可以清空所有单元格的运行结果再执行上述操作。
• 执行完上述步骤后可以用nvidia-smi命令查看显存占用情况，确认释放效果。

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