CUDA out of memory: 一种不常见bug的解决方法

问题描述

博主在跑推荐系统代码的时候出现了OOM(Out Of Memeory)错误，于是把device改为torch.device(“cpu”)，然而又出现了RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 10.91 GiB total capacity; 8.47 GiB already allocated; 1.75 MiB free; 10.22 GiB reserved in total by PyTorch)的报错，且报错定位在pickle.load()代码行处。。。

于是果断放弃使用cpu跑的想法，改用多卡并行加载数据，代码如下：

gpus = [0, 1]
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(gpu) for gpu in gpus)

结果同样的错误又出现了。。。

解决方案

既然错误出现在pickle.load()上，虽然不明白为什么报错，但也只能硬改了。用joblib.load()替换，还是不行，再用torch.save()和torch.load()将全局的pickle.dump()和pickle.load()替换，奇迹发生了！！！程序终于可以正常运行了！！

总结

造成爆显存的原因可能有多种，使用多张显卡对大多数这个问题来说都是一种较好的解决方法，博主所遇到的算是一个比较特殊的错误原因，希望可以给遇到类似问题的朋友一点参考。