当服务器中有多张显卡时可能会出现这个问题。
模型参数加载:

model_recover = torch.load(args.model_recover_path)

报错如下:

RuntimeError: CUDA error: out of memory

原因:
当使用torch.load加载模型参数时,会默认加载在第一块GPU0上,当GPU0没有在使用时,问题不大,但是显存被占满时这时候就加载不了了。
解决方法:

model_recover = torch.load(args.model_recover_path, map_location={'cuda:0': 'cuda:2'})

指定GPU映射,将GPU0映射到GPU2(任意一张空闲显卡),这样原本默认加载在GPU0上的参数将会加载到GPU2中,问题解决。

参考:torch代码运行时显存溢出问题之问题四

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐