ML-CUDA out of memory (torch.load())解决
1. 杂话
在上次的CUDA out of memory的问题解决之后(CUDA out of memory (Training)解决),这次又是遇到了新的问题。
使用Pytorch的朋友可能都知道,我们是可以使用torch.load()来加载我们之前训练好的模型的。
2. 问题
于是今天,我在使用torch.load()的时候就又碰到了爆显存的问题
但是我仔细看了一下我的代码,我在加载模型的时候,好像也没有用到那块被占用的显卡啊
(自从上次的报错之后,我和实验室的另一位同学达成了一人使用一块显卡的和平相处协议)
感谢这位大哥的帖子,帮我解决了问题pytorch加载模型torch.load报cuda out of memory解决方案
3. 解决
简单来说就是,在torch.save()得到的pth文件里,会保存着你之前是在那张显卡训练这个模型的信息,而你在加载这个模型的时候,就会经过这张显卡,这个时候,哪怕你最终不是使用这张显卡计算,但是这张显卡现在是被占用的话,就会报CUDA out of memory的错误。。。
也就是说,因为我和大哥一人一张显卡,之前我一直都是在cuda:1进行训练的,但是今天1卡被大哥用了,就算我把device改为cuda:0,在加载模型的时候也还是会报错,这个时候你就需要在代码中加入一个参数,来直接把模型映射到你要使用的显卡上
于是代码就从
torch.

本文介绍了在使用Pytorch时遇到CUDA内存不足问题的解决方法。当使用`torch.load()`加载模型时,由于.pth文件保存了训练时所用GPU的信息,即使更改设备设置,仍可能导致错误。解决方案是在加载模型时指定`map_location='cuda:0'`,将模型映射到目标GPU,从而避免错误。此方法帮助博主成功解决了跨GPU加载模型导致的内存冲突问题。
最低0.47元/天 解锁文章
2208





