CODEWITHSUNDEEP
pytorchhuggingface-transformers
yanachen
yanachen

Reputation: 3753

How to remove the model of transformers in GPU memory

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)

Hello, my codes can load the transformer model, for example, CTRL here, into the gpu memory. How to remove it from GPU after usage, to free more gpu memory?

show I use torch.cuda.empty_cache() ?

Thanks.

Upvotes: 5

Views: 13985

Answers (2)

ejb
ejb

Reputation: 86

@pierlj's solution doesn't seem to work for transformer models, however this solution from the thread they linked works for me:

import gc

del model
gc.collect()
torch.cuda.empty_cache()

Upvotes: 4

pierlj
pierlj

Reputation: 114

You can simply del tokenizer_ctrl and then use torch.cuda.empty_cache().

See this thread from pytorch forum discussing it.

Upvotes: 8

Related Questions