如何计算显存的占用,预防out of memory?
最近一次组会上,师兄点评一篇文章显存占用过多,突然发现还不知道如何具体的计算显存,只好去学习一下。
显存类似于内存,可以存放模型数据,参数等等;显存越大,所能运行的网络也就越大
torch.FatalError: cuda runtime error (2) : out of memory at /opt/conda/.......
out of memory: 显存装不下你那么多的模型权重还有中间变量
GPU计算单元用来进行数值计算,衡量计算量的单位是flop,浮点数先乘后加算一个flop计算能力越强大,速度越快。衡量计算能力的单位是 flops: 每秒能执行的 flop数量。
2*2+2 :1个flop
2*2+3*3+4*4 : 3个flop
1、 存储指标
1 Byte = 8 bit 1 K = 1024 Byte 1 M = 1024 K 1 G = 1024 M
除此之外,
1 Byte = 8 bit 1 KB = 1000 Byte 1 MB = 1000 KB 1 GB = 1000 MB 1TB = 1000 GB
常用的数值类型:
若一张256*256的RGB图片存储在显存中占有显存为(float):
3*256*256*4=0.75M,若batchsize=100,也就占用75M,显存,显然,占用显存较大的不是输入图片数据,那会是什么呢?
什么占用了显存?
首先,了解神经网络的构成,我们当然知道神经网络只是一种类似神经的架构,主要由构成网络层的各种参数构成,以及神经网络的各种中间输出。
网络模型的参数:
看一个例子:
- 模型权重:各种网络层的参数
- 卷积层,通常的conv2d
- 全连接层,也就是Linear层
- BatchNorm层
- Embedding层
- 中间变量:各种网络层的输出
而不占用显存的则是:
- 刚才说到的激活层Relu等
- 池化层
- Dropout层
具体计算方式:
- Conv2d(Cin, Cout, K): 参数数目:Cin × Cout × K × K
- Linear(M->N): 参数数目:M×N
- BatchNorm(N): 参数数目: 2N
- Embedding(N,W): 参数数目: N × W
参数占用显存:
参数占用显存 = 参数数目×n
n = 4 :float32 n = 2 : float16 n = 8 : double64
优化器的显存占用:
例如SGD优化器:
除了保存W之外还要保存参数对应的梯度,因此显存占用等于参数占用的显存的2倍。
Momentum-SGD:保存参数、梯度、动量------3倍
Adam:------------------------------------------4倍
输入输出的显存占用:
特点:
- 需要计算每一层的feature map的形状(多维数组的形状)
- 模型输出的显存占用与 batch size 成正比
- 需要保存输出对应的梯度用以反向传播(链式法则)
- 模型输出不需要存储相应的动量信息(因为不需要执行优化)
具体计算:
显存占用 = 模型显存占用 + batch_size × 每个样本的显存占用
注意 : 输入数据不用计算梯度;激活函数不用保存输入;
如何减小显存占用?出现out of memory如何处理?
- 尽量不使用全连接层
- 下采样
- 减小batchsize
最简单处理方法,也是最常用的方法
减小batchsize
一般模型参数与batchsize成一定的不严格的正比关系。
参考资料:https://blog.csdn.net/liusandian/article/details/79069926