完美解決pytorch载入预训练权重时出现的CUDA error: out of memory_cuda超出内存

今天小虎远程操纵工作站，想把昨晚练好的预训练模型迁移一下，发现跟往常不一样，nvidia的Cuda报错说出现超出内存的情况。

问题可能原因

我猜测可能是运行的进程没有完全结束（经常用ctrl + c打断进程），所以出现了这种错误。但是又不想重开机，因为还有其他程序在跑。有网友说用下面的指令彻底kill掉进程（这种操作很危险，容易直接把所有GPU的程序都打断）：

ps -elf | grep python

kill -9 [pid]
1
2
3

但是我并没有发现相关进程。
所以我试着用另一种办法：将权重load进入cpu，再让模型加载。避免load出现gpu有关的内存报错。

解决代码

我这里只给出我改动的部分，也是关键部分。使用GPU来load模型有可能导致在主机有多个GPU的环境下占用其他GPU内存，所以这里不如指定为cpu。
原始：

    if torch.cuda.is_available():
        PretrainedDict = torch.load(PreTrainedWeight)
    else:
        PretrainedDict = torch.load(PreTrainedWeight, map_location=torch.device('cpu'))
1
2
3
4

改动后直接用cpu载入，小虎尚未遇到直接用cpu载入权重有bug的情况：

    PretrainedDict = torch.load(PreTrainedWeight, map_location='cpu')
1

当然你也可以指定为你模型训练所用的GPU（如cuda:0），或者多GPU训练就不用特意定为哪个。

原始方案

改动后：

    if torch.cuda.is_available():
        try:
            PretrainedDict = torch.load(PreTrainedWeight)
        except:
            print("Cuda is out of memory")
        finally:
            PretrainedDict = torch.load(PreTrainedWeight, map_location=torch.device('cpu'))
    else:
        PretrainedDict = torch.load(PreTrainedWeight, map_location=torch.device('cpu'))
1
2
3
4
5
6
7
8
9

由于这种改动仍然会存在占用其他GPU问题，所以弃用，2023/04/28。

参考资料

Out of memory error when resume training even though my GPU is empty
CUDA error: out of memory when load models
CUDA error when loading my model