torch有时候跑着跑着显存吃满了，就会报错：RuntimeError: CUDA out of memory. Tried to allocate 916.00 MiB (GPU 0; 6.00 GiB total capacity; 4.47 GiB already allocated; 186.44 MiB free; 4.47 GiB reserved in total by PyTorch)

这种情况怎么办呢？总结一下排查错误的方法。

1. 看一下显存是不是真的满了。有可能机子上有多张卡，用nvidia-smi看一下有几张卡。这个指令只能输出当前时刻的显存占用。所以题主一般用动态追踪的方式。输入指令：watch -n 1 nvidia-smi这样终端里就会每一秒更新一下显卡状态，如果想更快一点也可以把1换成0.5之类的。输出如下所示。可以看到这个机子有两张A5000，第一张卡显存是23953MiB / 24564MiB，快用完了，第二张是 18372MiB / 24564MiB，还有一点可以用。

   +-----------------------------------------------------------------------------+
   | NVIDIA-SMI 515.65.01    Driver Version: 515.65.01    CUDA Version: 11.7     |
   |-------------------------------+----------------------+----------------------+
   | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
   | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
   |                               |                      |               MIG M. |
   |===============================+======================+======================|
   |   0  NVIDIA RTX A5000    Off  | 00000000:17:00.0 Off |                  Off |
   | 51%   67C    P2   111W / 230W |  23953MiB / 24564MiB |     68%      Default |
   |                               |                      |                  N/A |
   +-------------------------------+----------------------+----------------------+
   |   1  NVIDIA RTX A5000    Off  | 00000000:65:00.0 Off |                  Off |
   | 60%   78C    P2   106W / 230W |  18372MiB / 24564MiB |     27%      Default |
   |                               |                      |                  N/A |
   +-------------------------------+----------------------+----------------------+
   
   +-----------------------------------------------------------------------------+
   | Processes:                                                                  |
   |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
   |        ID   ID                                                   Usage      |
   |=============================================================================|
   |    0   N/A  N/A     13291      G   /usr/lib/xorg/Xorg                  4MiB |
   |    0   N/A  N/A     13998      G   /usr/lib/xorg/Xorg                  4MiB |
   |    0   N/A  N/A     20176      C   python                           7977MiB |
   |    0   N/A  N/A     20407      C   python                           7977MiB |
   |    0   N/A  N/A     20693      C   python                           7977MiB |
   |    0   N/A  N/A     25065      G   /usr/lib/xorg/Xorg                  4MiB |
   |    0   N/A  N/A     25369      G   /usr/lib/xorg/Xorg                  4MiB |
   |    1   N/A  N/A     13291      G   /usr/lib/xorg/Xorg                 15MiB |
   |    1   N/A  N/A     13386      G   /usr/bin/gnome-shell              108MiB |
   |    1   N/A  N/A     13998      G   /usr/lib/xorg/Xorg                194MiB |
   |    1   N/A  N/A     14085      G   /usr/bin/gnome-shell               82MiB |
   |    1   N/A  N/A     14979      G   /usr/lib/firefox/firefox            6MiB |
   |    1   N/A  N/A     19570      C   python                           1849MiB |
   |    1   N/A  N/A     20200      G   gnome-control-center               27MiB |
   |    1   N/A  N/A     20830      C   python                           7977MiB |
   |    1   N/A  N/A     23158      C   python                           7977MiB |
   |    1   N/A  N/A     25065      G   /usr/lib/xorg/Xorg                 39MiB |
   |    1   N/A  N/A     25090      G   /usr/bin/gnome-shell               12MiB |
   |    1   N/A  N/A     25369      G   /usr/lib/xorg/Xorg                 16MiB |
   |    1   N/A  N/A     25455      G   /usr/bin/gnome-shell               58MiB |
   +-----------------------------------------------------------------------------+
   

2. 既然第二张卡还剩一些显存，为什么跑代码后还是报错RuntimeError: CUDA out of memory.呢？首先要确保你的代码跑起来之后是第二张卡能塞下的。比如我的程序只占用2000MiB，理论上第二张卡一定是可以塞下的，这个时候报错那就有可能默认放到第一张卡上去了。解决方法是指定用一张卡，一般我会在命令行指定卡，而不是代码里。在原本的命令前加上CUDA_VISIBLE_DEVICES=1就好了。

   CUDA_VISIBLE_DEVICES=1  python main.py
   

3. 如果上面都检查过了，那有可能就是代码问题了。我debug的思路是每次注释掉一部分代码，再用watch -n 1 nvidia-smi这个指令观察显存有没有爆炸，以此定位导致显存爆炸的代码。代码问题就多种多样了，举一些常见例子。

   • 3.1 batch_size太大了。到了要检查代码这一步首先是看batch_size，但如果batch_size太大了应该是一个epoch都跑不了，在前向传播的时候甚至前向传播之前就会出问题了。所以如果代码是一次前向传播都无法完成，应该首先排查batch_size.

   • 3.2 优化器没有step。这会导致training跑了一些iteration以后才显存爆炸。optimizer.step()这行如果漏掉了的话梯度会一直累积，显存一直不被释放。

   • 3.3 把一个cuda上的tensor和一个常量做运算。这种情况一般也是会导致跑了一些iteration以后才显存爆炸。两个不同设备上的tensor运算，这个错误一般torch会报错的，但是也有非常隐晦的情况，比如不指明在哪个设备上的常量，需要手动发现。你得到的结果其实是在cuda上的，如：

     >>> a = torch.Tensor(1).cuda()	
     	# 顺便一提 如果要给a赋值1应该是a = torch.Tensor([1]).cuda()，不加中括号赋值是错的
     >>> a == 1
     tensor([False], device='cuda:0')
     

     这个时候应该妥善处理结果！如果得到的tensor没有被赋给一个变量，可能会一直放在显存里无法释放。我一般直接移到cpu上运算。

   • 3.4 需要存下一些cuda tensor的情况。这种情况一般会导致显存占用比较大，有可能导致显存爆炸。比如：

      outs = [out.cpu() for out in outs]
      all_outs = torch.cat([all_outs.cpu(),torch.stack(outs).cpu()],dim=1)
     

     这种情况我明明已经把outs这个list里面的每个tensor全部移到cpu，还是会占用很大显存。因为这是在inference时候的代码，我后续不需要它的梯度，解决方法是把梯度detach掉：

     outs = [out.detach().cpu() for out in outs]
     all_outs = torch.cat([all_outs.cpu(),torch.stack(outs).cpu()],dim=1)
     

     这样就不占显存啦～～