今天跑一個模型的時候，需要加載部分預訓練模型的參數，這期間遇到使用torch.load 忽略了 map_location參數 默認gpu，這導致這個變量分配的顯存 不釋放 然后占用大量資源 gpu資源不能很好的利用。

問題講解：

比如我們一般我們會使用下面方式進行加載預訓練參數 到 自身寫的模型中：

from transformers import RobertaForMultipleChoice
import torch
model = RobertaForMultipleChoice.from_pretrained("roberta-large")
pretrained_model = torch.load("./checkpoints/txt_matching_e1.pth").roberta
pretrained_dict = pretrained_model.state_dict()
model_dict = model.roberta.state_dict()
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} #去除一些不需要的參數
model_dict.update(pretrained_dict)
model.roberta.load_state_dict(model_dict)

1. 當我們沒有使用參數時候 load 默認使用了一塊顯卡然后報錯

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 3; 10.76 GiB total capacity; 350.54 MiB already allocated; 21.81 MiB free; 356.00 MiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

torch load 之前gpu使用

torch load 之后 outof memory 了 并且也不釋放

2. 當我們沒有使用參數時候 load 默認使用了一塊顯卡然后報錯
當我試試指定顯卡 gpu會使用2841
pretrained_model = torch.load(“./checkpoints/txt_matching_e1.pth”,map_location=‘cuda:0’).roberta!

(model 直接cuda 的gpu 占用情況)
然后把這里面參數給model，并且model也是用cuda0 然后gpu使用4193

你可能會想model是不是model load預訓練參數之后 就這么大了 那么load 和 load 參數后model 用不同gpu看看。

原理：cuda的內存管理機制
參考解釋博客：Pytorch訓練模型時如何釋放GPU顯存

解決方案：

1. 不占用顯存的使用方法，使用cpu 然后在del 用gc釋放內存

model = RobertaForMultipleChoice.from_pretrained("roberta-large")
pretrained_model = torch.load("./checkpoints/txt_matching_e1.pth",map_location='cpu').roberta
pretrained_dict = pretrained_model.state_dict()
model_dict = model.roberta.state_dict()
model_dict.update(pretrained_dict)
model.roberta.load_state_dict(model_dict)
del pretrained_model
import gc
gc.collect()

2. 合理使用， torch.cuda.empty_cache()

這個需要了解一下python的內存管理，引用機制。
比如我pretrain_model 給model直接加載參數，model和pretrain_model 都在cuda:0上，使用torch.cuda.empty_cache() 不能釋放pretrain_model 的顯存。
當 我把model 放到 cuda：1上（本來在cuda:0）,這時候用torch.cuda.empty_cache() 可以釋放。