如何去除 GPU 内存中的变压器模型

Posted 2023-03-29

【中文标题】如何去除 GPU 内存中的变压器模型

【英文标题】：How to remove the model of transformers in GPU memory

【发布时间】：2021-11-20 06:26:41

【问题描述】：

from transformers import CTRLTokenizer, TFCTRLLMHeadModel
tokenizer_ctrl = CTRLTokenizer.from_pretrained('ctrl', cache_dir='./cache', local_files_only=True)
model_ctrl = TFCTRLLMHeadModel.from_pretrained('ctrl', cache_dir='./cache', local_files_only=True)
print(tokenizer_ctrl)
gen_nlp  = pipeline("text-generation", model=model_ctrl, tokenizer=tokenizer_ctrl, device=1, return_full_text=False)

您好，我的代码可以将transformer模型，例如这里的CTRL，加载到gpu内存中。 使用后如何将其从 GPU 中移除，以释放更多 gpu 内存？

显示我使用torch.cuda.empty_cache()？

谢谢。

【参考方案1】：

您可以简单地del tokenizer_ctrl，然后使用torch.cuda.empty_cache()。

看到这个来自 pytorch 论坛的 thread 讨论它。