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

Question

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)

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

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

谢谢。

Answer 1

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

请看pytorch论坛上讨论它的thread。