解决加载torch模型时出现CUDA out of memory_torch.cuda.outofmemoryerror: cuda out of memory.

作者：Monodyee | 2024-02-16 00:48:26

踩

torch.cuda.outofmemoryerror: cuda out of memory.

解决加载torch模型时出现CUDA out of memory

正常来说出现“CUDA out fo memory”是CUDA内存不够出现的bug。
事情是这样滴，我训练完一个模型之后，加载的时候并没有把模型加载到gpu，但是还是报错，代码如下：

from transformers import BertForSequenceClassification, AdamW

pretrained = 'bert-base-chinese'

model = BertForSequenceClassification.from_pretrained(
    pretrained, 
    num_labels=121, 
    output_attentions = False, # 模型是否返回 attentions weights.
    output_hidden_states = False, # 模型是否返回所有隐层状态.
    return_dict=False
)
model.load_state_dict(torch.load("model/BERT_model.h5"))
1
2
3
4
5
6
7
8
9
10
11
12

在没有加载gpu的情况下为什么还会爆出CUDA内存不够的bug呢？
————————————————————————————
原因
保存时使用了某个gpu，加载的时候自动会经过这个gpu，但是这时候如果这个gpu内存不够的话就会报错
解决方法
在load_state_dict加一个变量

model.load_state_dict(torch.load("model/BERT_model.h5", map_location='cpu'))
1

这样就不会报CUDA内存不够的bug了，问题解决。
（ps：大家如果没有跑模型的时候记得把内存释放掉）

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/Monodyee/article/detail/89004