我们平时用 PyTorch 训练模型，有时会碰到显存不足的情况。除了找对应的解决办法，比如累加梯度、使用自动混合精度，我们还应该弄清楚，训练时显存究竟在哪些环节被大量占用。主要有下面这几个环节：

1. CUDA 运行内存
2. 模型的固定参数
3. 模型的前向计算
4. 模型的反向计算
5. 优化方法统计量

1. 运行内存

CUDA（Compute Unified Device Architecture，计算统一设备架构）是显卡厂商 NVIDIA 推出的运算平台。通过它，我们才能调用 GPU 的处理能力，大幅提升计算性能。

对我们来说，CUDA 本质上是一套跑在 GPU 硬件上的软件平台，我们的计算任务要建立在这个平台之上，才能用上 GPU 的算力。既然它本身也是一套软件，运行起来自然会占一部分显存，占多少得看 CUDA 的版本，有的占 600M 左右，有的会占到 1G 以上。

在说 CUDA 占多少之前，我们先了解下 PyTorch 的显存使用机制。GPU 显存就是我们手头全部可用的资源。写过 C/C++ 的同学都知道，频繁地申请和释放资源，比如 C 的 malloc/free、C++ 的 new/delete，会明显拖慢系统性能。为了减少这类操作，就有了资源池的概念。它的思路是，预先从全部可用资源里申请一大块，程序需要资源时就从池子里拿，省掉复杂耗时的系统调用，用完再放回池子里，池子不够了再向系统要。这样一来，资源使用这个环节的效率就高了。

PyTorch 给张量分配显存也是这套做法，先申请一大块内存，张量要用时从内存池里取，不用了就归还。可以想见，如果没有这种资源缓存机制，运行效率会非常慢。

接下来，我们用一段代码来验证 CUDA 平台一运行就会占用部分显存。先装一个库：

pip install pynvml

import torch
import pynvml


# 初始化 pynvml 库
pynvml.nvmlInit()
convert = lambda x: int(x / 1024 / 1024)
# 获得显卡设备对象
device_object = pynvml.nvmlDeviceGetHandleByIndex(0)


# 查看显存资源
def show_usage():
    # 获得显存信息
    device_memory = pynvml.nvmlDeviceGetMemoryInfo(device_object)
    # 全部可用显存
    total = convert(device_memory.total)
    # 已经使用显存
    used = convert(device_memory.used)
    # 剩余可用显存
    free = convert(device_memory.free)
    print('总共:', total, '使用:', used, '剩余:', free)


# 1. CUDA 初始化会占用部分显存
def test01():
    show_usage()
    # 如果张量创建在 CPU 是不会占用显存，并且也不会初始化 CUDA
    torch.tensor(0.0, device='cpu')
    show_usage()
    torch.tensor(0.0, device='cuda')
    # 清空缓存
    torch.cuda.empty_cache()
    show_usage()

if __name__ == '__main__':
    test01()

程序输出结果：

总共: 5932 使用: 0 剩余: 5932
总共: 5932 使用: 0 剩余: 5932
总共: 5932 使用: 586 剩余: 5346

上面这段代码如果不清空缓存，输出会是 588 而不是 586，即 588 等于 586 加上 PyTorch 自己缓存的部分。另外，我们创建的 cuda 张量没有被任何变量引用，创建完就会被自动回收，这时候清空缓存才得到 586，否则仍是 588。原因在于，每次向 cuda 设备创建张量，分配的显存都是 512 的倍数。

import torch

def test02():

    # 0
    print(torch.cuda.memory_allocated())
    a = torch.tensor(0.0, device='cuda')
    # 512
    print(torch.cuda.memory_allocated())
    # 1024
    b = torch.tensor(0.0, device='cuda')
    print(torch.cuda.memory_allocated())
    # 1536
    c = torch.tensor(0.0, device='cuda')
    print(torch.cuda.memory_allocated())



if __name__ == '__main__':
    test02()

程序输出结果：

0
512
1024
1536

torch.cuda.memory_allocated 可以获得目前分配的内存数量。

2. 固定参数

这一部分就比较好理解了，加载模型本质上就是加载模型参数，所以参数本身会占一部分显存。默认情况下，PyTorch 里的参数用的是 float32 类型。我们看下面的代码：

import torch
import torch.nn as nn


def test01():
    print(torch.cuda.memory_allocated())
    linear = nn.Linear(in_features=1, out_features=1, bias=False).cuda()
    print(torch.cuda.memory_allocated())

if __name__ == '__main__':
    test01()

程序输出结果：

0
512

前面创建的线性层不带偏置，只有一个参数，按说只该占 4 字节，为什么这里显示 512 字节呢？原因是 PyTorch 分配显存时按 512 的倍数分配，也就是按块分。这样不怕浪费吗？这也是出于效率考虑，按块分配便于内存管理，尽量避免产生内存碎片。

import torch
import torch.nn as nn

def test02():
    print(torch.cuda.memory_allocated())
    linear = nn.Linear(in_features=128, out_features=1, bias=False).cuda()
    print(torch.cuda.memory_allocated())


if __name__ == '__main__':
    test02()

输出仍然是 512 字节。如果把 in_features 从 128 换成 129，就会分配 1024 字节显存。注意一个参数的大小是 4 字节。

思考：下面的模型占用多大显存？

import torch
import torch.nn as nn


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.linear1 = nn.Linear(1, 1, bias=False)
        self.linear2 = nn.Linear(1, 1, bias=False)

    def forward(self, inputs):
        inputs = self.linear1(inputs)
        inputs = self.linear2(inputs)
        return inputs


def test03():
    print(torch.cuda.memory_allocated())
    model = Net().cuda()
    print(torch.cuda.memory_allocated())


if __name__ == '__main__':
    test03()

程序输出结果是：

0
1024

3. 前向反向计算

网络做前向计算时会把中间结果保存下来。为什么要保存？因为反向求梯度时要用到这些中间结果。而反向算出来的梯度值本身也需要显存来存，所以前向和反向计算都会占显存。

另外，输入的 batch_size 越大，占用的显存也越大。

import torch
import torch.nn as nn


def test():

    print(torch.cuda.memory_allocated())

    model = nn.Linear(1, 1).cuda()
    print(torch.cuda.memory_allocated())

    # 前向计算
    # 5120 = 1024 + 4096（1024 个输入大小）
    inputs = torch.randn(size=(1024, 1)).cuda()
    print(torch.cuda.memory_allocated())

    # 正向计算需要缓存中间计算结果(outputs)
    # 注意：用变量承接相当于缓存了中间结果
    # 9216 = 5120 + 4096（1024 个缓存结果）
    outputs = model(inputs)
    print(torch.cuda.memory_allocated())

    # 计算损失
    # 9728 = 9216 + 512 缓存损失结果
    loss = torch.mean(outputs)
    print(torch.cuda.memory_allocated())

    # 反向计算
    # 10752 = 9728 + 512 保存梯度值
    loss.backward()
    print(torch.cuda.memory_allocated())


if __name__ == '__main__':
    test()

程序执行结果：

0
1024
5120
9216
9728
10752

反向传播之后，可以释放 outputs、loss 这些变量。

4. 优化统计量

不同的优化器内部会维护一些统计量。比如 SGD 会记录每个参数的历史动量梯度，Adam 会记录每个参数的一阶、二阶梯矩估计。这些在训练过程中同样要占显存，而且参数量越大，优化器占的显存就越大。

import torch
import torch.nn as nn
import torch.optim as optim


def test():

    # 0
    print(torch.cuda.memory_allocated())

    # 512
    model = nn.Linear(1, 1, bias=False).cuda()
    print(torch.cuda.memory_allocated())

    # 1024
    inputs = torch.randn(size=(1, 1)).cuda()
    print(torch.cuda.memory_allocated())

    # 1536
    outputs = model(inputs)
    print(torch.cuda.memory_allocated())

    # 2048
    loss = torch.mean(outputs)
    print(torch.cuda.memory_allocated())

    # 2560
    loss.backward()
    print(torch.cuda.memory_allocated())

    # 3584
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    optimizer.step()
    print(torch.cuda.memory_allocated())


if __name__ == '__main__':
    test()

程序执行结果：

0
512
1024
1536
2048
2560
3584

SGD 一旦设置了 momentum，内部就会为每个参数记录一份历史梯度，Adam 要记录的东西更多，所以 Adam 的显存占用通常比 SGD 更大。