pytorch在训练时如何优化显存?显存不够用了

img


太无奈了,不知道该咋办?有什么方法可以优化吗?在哪里可以省显存?

1.计算总loss时候进行代码优化
loss本身是一个包含梯度信息的 tensor,正确的求损失和的方式为:

total_loss += loss.item()
1
2.降低batch_size
适当降低batch size, 则模型每层的输入输出就会成线性减少, 效果相当明显。必须要多个batch_size时候可以采用梯度累计的方式,但是这比上直接降低batch_size占用显存会大点。

传统方式:

for i,(feature,target) in enumerate(train_loader):
    outputs = model(feature)  # 前向传播
    loss = criterion(outputs,target)  # 计算损失

    optimizer.zero_grad()   # 清空梯度
    loss.backward()  # 计算梯度
    optimizer.step()  # 反向传播, 更新网络参数

梯度累计:

for i,(features,target) in enumerate(train_loader):
    outputs = model(images)  # 前向传播
    loss = criterion(outputs,target)  # 计算损失
    loss = loss/accumulation_steps   # 可选,如果损失要在训练样本上取平均

    loss.backward()  # 计算梯度
    if((i+1)%accumulation_steps)==0:
        optimizer.step()        # 反向传播,更新网络参数
        optimizer.zero_grad()   # 清空梯度

3.Relu 的 inplace 参数
激活函数 Relu() 有一个默认参数 inplace ,默认为Flase, 当设置为True的时候,我们在通过relu() 计算得到的新值不会占用新的空间而是直接覆盖原来的值,这表示设为True, 可以节省一部分显存。

4.释放不需要的张量和变量
采用del释放你不再需要的张量和变量,这也要求我们在写模型的时候注意变量的使用,不要随心所欲,漫天飞舞。

5.数据变小
图像适当的缩小可以极大的降低参数量。

6.精简模型
减少卷积核数量; 尽量少的使用 Linear 等。全连接层参数较多,较少参数或则不用全连接层。使用全局平均池化进行替代。

深度学习显存占用取决于:

1、模型所有带参数权重的层的输入输出(在梯度更新时候要计算梯度)

2、模型自身的参数大小,因为梯度下降需要用到梯度,大小2;

3,有些优化器需要用到一阶二阶动量,这些动量需要用到原参数大小的权重,所以在adam优化器中,参数占用需要*4。

显存的优化技巧:

1、数据流使用生成器,比如tensorflow的tfrecords和pytorch的dataloder,也可以自建生成器。

2、对于没有参数的层可以使用inplace操作,如在pytorch中,relu激活函数就有inplace选项。假设x是输入,y是输出,这个操作的含义是在计算出y之后,用y代替x。在进行梯度更新的时候,用y的值计算出x进行梯度更新。这个方法比较适用于简单的无参数层,在过于复杂的情况下,等于牺牲时间成本换空间。

3、梯度累计方法,将一个完整batch size分为多份,每次计算后不更新梯度,全部计算完成才进行更新。

4、选择比较轻量级别的优化器。

5、apex半精度。

5、在自然语言处理任务中,padding的处理非常重要。无论是CNN,RNN,我们都应当使用动态的结构(动态rnn,避免静态RNN)来避免padding部分的计算。在有些nlp任务中,不是每个单词的表征层都要被用来predict,我们可以只选择有用的词来进行最后预测层的计算。

  1. 把bitch_size的大小调小点,2的倍数;
  2. num_work改成1或者0,减少同时进程数;

1.减少batch size(最直观有效的做法)
2.通过half()函数转为半精度,由float32改为float16,建议去查一下Pytorch半精度模型。

显存不够用,可以很多方式处理:

  • 物理处理,增加自己的显存;
  • 训练参数设置,较小batch size;
  • 数据大小的处理,有很多时候,并不需要那么大的数据,可以将数据减小一半,显存需求就会减小一半,也很有效;
  • 数据格式调整,将float32, 调整为float16,甚至调整为float8;
  • 程序优化,在显卡运算的时候,避免中间变量,比如a, b, c三个变量,如果三者有某种关系,可以改写成两个,甚至一个变量,如果中间变量只是变量而已,可以使用del a来节省;
  • 代码内部优化,可以看这文章:https://blog.csdn.net/weixin_43914632/article/details/120759812

torch.FatalError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THC/generic/THCStorage.cu:58 显存不够了。。。
如何计算
首先我们应该了解一下基本的数据量信息:

1 G = 1000 MB
1 M = 1000 KB
1 K = 1000 Byte
1 B = 8 bit
好,肯定有人会问为什么是1000而不是1024,这里不过多讨论,只能说两种说法都是正确的,只是应用场景略有不同。这里统一按照上面的标准进行计算。
然后我们说一下我们平常使用的向量所占的空间大小,以Pytorch官方的数据格式为例(所有的深度学习框架数据格式都遵循同一个标准):

img


一般一个8-bit的整型变量所占的空间为1B也就是8bit。而32位的float则占4B也就是32bit。而双精度浮点型double和长整型long在平常的训练中我们一般不会使用。

ps:消费级显卡对单精度计算有优化,服务器级别显卡对双精度计算有优化。

也就是说,假设有一幅RGB三通道真彩色图片,长宽分别为500 x 500,数据类型为单精度浮点型,那么这张图所占的显存的大小为:500 x 500 x 3 x 4B = 3M。
而一个(256,3,100,100)-(N,C,H,W)的FloatTensor所占的空间为256 x 3 x 100 x 100 x 4B = 31M
看起来一张图片(3x256x256)和卷积层(256x100x100)所占的空间并不大,那为什么我们的显存依旧还是用的比较多,原因很简单,占用显存比较多空间的并不是我们输入图像,而是神经网络中的中间变量以及使用optimizer算法时产生的巨量的中间参数。

我们首先来简单计算一下Vgg16这个net需要占用的显存:
通常一个模型占用的显存也就是两部分:
模型自身的参数(params)
模型计算产生的中间变量(memory)

img


图片来自cs231n,这是一个典型的sequential-net,自上而下很顺畅,我们可以看到我们输入的是一张224x224x3的三通道图像,可以看到一张图像只占用150x4k,但上面标注的是150k,这是因为上图中在计算的时候默认的数据格式是8-bit而不是32-bit,所以最后的结果要乘上一个4。
我们可以看到,左边的memory值代表:图像输入进去,图片以及所产生的中间卷积层所占的空间。我们都知道,这些形形色色的深层卷积层也就是深度神经网络进行“思考”的过程:

img


图片从3通道变为64 –> 128 –> 256 –> 512 …. 这些都是卷积层,而我们的显存也主要是他们占用了。

还有上面右边的params,这些是神经网络的权重大小,可以看到第一层卷积是3×3,而输入图像的通道是3,输出通道是64,所以很显然,第一个卷积层权重所占的空间是 (3 x 3 x 3) x 64。
另外还有一个需要注意的是中间变量在backward的时候会翻倍!
要注意,优化器也会占用我们的显存!
为什么,看这个式子:

img


上式是典型的SGD随机下降法的总体公式,权重W在进行更新的时候,会产生保存中间变量
模型中哪些层会占用显存
有参数的层即会占用显存的层。我们一般的卷积层都会占用显存,而我们经常使用的激活层Relu没有参数就不会占用了。

占用显存的层一般是:

卷积层,通常的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
额外的显存
总结一下,我们在总体的训练中,占用显存大概分以下几类:

模型中的参数(卷积层或其他有参数的层)
模型在计算时产生的中间参数(也就是输入图像在计算时每一层产生的输入和输出)
backward的时候产生的额外的中间参数
优化器在优化时产生的额外的模型参数
但其实,我们占用的显存空间为什么比我们理论计算的还要大,原因大概是因为深度学习框架一些额外的开销吧,不过如果通过上面公式,理论计算出来的显存和实际不会差太多的。
如何优化
优化除了算法层的优化,最基本的优化无非也就一下几点:
减少输入图像的尺寸
减少batch,减少每次的输入图像数量
多使用下采样,池化层
一些神经网络层可以进行小优化,利用relu层中设置inplace
购买显存更大的显卡
从深度学习框架上面进行优化

train.py里batch size(批次大小)调小点,img_size调小一点
预训练模型换内存小点的,或者不使用
更换显卡或者每次训练完成后重启一下,清空无用显存
每次GPU读取图像的数量设置小一点,应该是per_img_GPU,调小点,程序里ctrl+F搜一搜
优化模型,减少卷积层数什么的,或者是dropout调大点试试

减小网络减小batch最有效