Pytorch 0号卡占用问题

RuntimeError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1587428266983/work/aten/src/THC/THCCachingHostAllocator.cpp:278

如上，我们在跑程序时经常会遇到这种情况，除了常规的因为模型参数量或计算量过大导致的问题，还有一种经常的情况下是因为0号显卡显存不足，导致显存空间分配不足，从而提示报错。很多人有一个疑问就是，为什么我们明明指定的是其它张卡，可0卡经常会多多少少占用一点，如果此时刚好有人在占用0号卡把显存给挤没了，即使你指定的卡不是0号卡，也会报“out of memory”。

解决这个问题，我们首先要先了解为什么会出现此问题。在Pytorch框架中，模型初始化的时候会默认占用0号卡的显卡，此时如果0号卡的显存空间不足，则会报错，具体讨论可见此链接。

因此，为了解决这个问题，最直接的方法便是将其它显卡“屏蔽”掉即可，换句话说就是使指定的卡可见，常用的代码有：

os.environ['CUDA_VISIBLE_DEVICES'] = ‘gpu_id’
device = torch.device(f'cuda:{0}' if torch.cuda.is_available() else 'cpu')

其中，gpu_id便是我们想要指定的设备号，多卡同理。或者直接在运行命令行的时候加入也可：

CUDA_VISIBLE_DEVICES=gpu_id python train.py

此外，还有一种方法是使用“with”语句指定，在涉及到调用显卡的地方指定设备号：

with torch.cuda.device(gpu_id):
	...

如果以上方法都不奏效，那么可以再试试看把模型map到cpu上运行，具体问题具体分析。