笔记本内容

谈谈显存 / 内存溢出 #

目前有很多用户都深受内存 / 显存溢出的困扰，本文就针对这个点来深度探讨（加实例演示）一下如何防止内存泄漏。

什么是显存 / 内存溢出？ #

内存溢出，通常也称为 OOM（Out Of Memory 的缩写），是编程中经常遇到的问题，特别是使用 CUDA 的深度学习训练程序更容易遇到，因为内存、显存都有可能溢出。

原理非常简单，因为机器的内存/显存是有上限的，因此当计算机的所有程序使用的内存或显存之和超过这个上限，就会导致程序报错或崩溃，甚至可能导致机器崩溃掉。

溢出的常见场景以及应对方法 #

因为显存和内存在功能上都是用于存储程序运行时的状态的，所以他们的溢出方式也大致相似。

一、一次性分配过多资源 #

这是最简单也是最容易发生的场景，特别是针对显存，因为显存一般比内存小得多。

遇到这种情况一般来说需要调整自己的代码参数（比如降低 Batch Size，使用更小的输入，或换用更小的模型）。

运行下面的代码后程序会报错，因为显存溢出了：

import torch

torch.rand(1024, 1024, 512, 512, 512).double().cuda()

只需要将上面代码中的 cuda() 去掉，就会报内存溢出的错误。

内存溢出一般来说不会报错，而是程序会直接被操作系统鲨掉，在 JupyterLab 中的表现就是 Kernel 直接停止，需要重新运行。这里 PyTorch 会报错是因为 PyTorch 本身做了保护。

二、有其他进程占用资源 #

有同学经常问这种问题：「为什么我只用了很小的模型，很小的 batch size 就报显存溢出了」，特别是对于使用 JupyterLab 的场景，特别容易出现此类问题。

如果遇到这种问题，有可能是之前启用的进程没有被完全鲨掉造成的。可以通过以下方式来找到占用资源的进程：

1. 查找占用显存的进程：

!nvidia-smi

Sun Sep 19 09:40:21 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.57.02    Driver Version: 470.57.02    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| 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 GeForce ...  Off  | 00000000:06:00.0 Off |                  N/A |
| 30%   20C    P8    15W / 320W |   3251MiB / 10018MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

为什么显存明明被占用了，但下面的 Processes 是空的？这是因为我的这个笔记本使用了 Docker 容器，目前 nvidia-smi 命令还不支持在容器中获得进程信息。这也是为什么在使用一些云计算平台的机器时，无法用此命令查看到占用显存的进程号的原因。

2. 查找占用内存的进程：

!ps -eo pid,cmd,%mem,%cpu --sort=-%mem | head -n 4

    PID CMD                         %MEM %CPU
    318 /environment/python/version 15.9  1.5
     31 /environment/python/version  0.2  1.5
     43 /environment/python/version  0.2  0.0

上面显示了当前内存占用 Top 3 的进程。

拿到进程号后，我们可以使用下面的命令来手动鲨掉进程：

kill <pid>

将换成进程号即可。

在 JupyterLab 中，可以通过左侧的 Kernel Tab 来关闭不需要的 Kernels（每一个 Kernel 都是一个 Notebook 进程）。

三、内存泄漏 #

内存泄漏使指程序在运行过程中不断的申请新的内存，而不释放原有的内存（原来的内存可能本应该释放）。程序在刚开始运行的时候不会有任何问题，但随着运行的时间增加，进程消耗的内存不断上涨，最终吃掉所有机器的内存而导致程序挂掉。

这是最难排查场景，属于程序本身的问题，多发生于很大的循环结构中（例如训练循环或验证循环）。

下面就是最简单的内存泄漏的例子：

result = []

for epoch in range(100):
    for iteration in range(10000):
        # 训练的代码
        predict_mask = torch.rand(4, 3, 512, 512)
        result.append(predict_mask)

执行代码后，程序将在运行一段时间后崩溃，能运行的时间取决于你的机器内存大小。JupyterLab 可能会发出如下的警告：

问题的原因是训练代码在每个迭代都会往 result 这个列表中添加对象，而这个对象本身也比较大，是一个 4 x 3 x 512 x 512 的张量，所以程序在跑到一定程度后因为 result 变量过大而导致内存耗尽，程序直接崩溃。

排查此类问题可以借助工具，在程序的不同阶段中打印一些 DEBUG 信息，然后同时使用 watch -n 0.2 free -m 命令来观察内存的变化情况。这样就能大致得出内存泄漏的点，然后再仔细排查。

笔记本

聊聊内存 / 显存溢出