【深度学习】为什么深度学习需要大内存?

【深度学习】为什么深度学习需要大内存?

本文主要译介自Graphcore在2017年1月的这篇博客: Why is so much memory needed for deep neural networks。介绍了深度学习中内存的开销,以及降低内存需求的几种解决方案。 为便于阅读,本文修改了原文分段,并添加更详细的计算说明。

深度学习的内存消耗在哪里?

回顾:简单例子

考虑一个单层线性网络,附带一个激活函数:

     h


     =



      w


      1



     x


     +



      w


      2




    h=w_1x+w_2


 h=w1​x+w2​






     y


     =


     f


     (


     h


     )



    y=f(h)


 y=f(h)

代价函数:

    E


    =


    ∣


    ∣


    y


    −



     y


     ‾



    ∣



     ∣


     2




   E=||y-\overline{y}||^2


E=∣∣y−y​∣∣2

在训练时,每一个迭代要记录以下数据:

  • 当前模型参数

          w
    
    
          1
    
    
    
         ,
    
    
    
          w
    
    
          2
    
    
    
    
        w_1,w_2
    
    
     w1​,w2​</li><li>前向运算各层响应:
    
    
    
    
         x
    
    
         ,
    
    
         h
    
    
         ,
    
    
         y
    
    
    
        x, h, y
    
    
     x,h,y</li>
    

    这样,可以在后向运算中用梯度下降更新参数:

         Δ
    
    
    
          w
    
    
          1
    
    
    
         =
    
    
         η
    
    
         ⋅
    
    
    
    
           ∂
    
    
           E
    
    
    
    
           ∂
    
    
    
            w
    
    
            1
    
    
    
    
    
         =
    
    
         η
    
    
         ⋅
    
    
         2
    
    
         (
    
    
         y
    
    
         −
    
    
    
          y
    
    
          ‾
    
    
    
         )
    
    
         ⋅
    
    
    
          f
    
    
          ′
    
    
    
         (
    
    
         h
    
    
         )
    
    
         ⋅
    
    
         x
    
    
    
        \Delta w_1=\eta\cdot \frac{\partial E}{\partial w_1}=\eta \cdot 2(y-\overline{y})\cdot f'(h) \cdot x
    
    
     Δw1​=η⋅∂w1​∂E​=η⋅2(y−y​)⋅f′(h)⋅x
    
    
    
    
    
    
         Δ
    
    
    
          w
    
    
          2
    
    
    
         =
    
    
         η
    
    
         ⋅
    
    
    
    
           ∂
    
    
           E
    
    
    
    
           ∂
    
    
    
            w
    
    
            1
    
    
    
    
    
         =
    
    
         η
    
    
         ⋅
    
    
         2
    
    
         (
    
    
         y
    
    
         −
    
    
    
          y
    
    
          ‾
    
    
    
         )
    
    
         ⋅
    
    
    
          f
    
    
          ′
    
    
    
         (
    
    
         h
    
    
         )
    
    
    
        \Delta w_2=\eta\cdot \frac{\partial E}{\partial w_1}=\eta \cdot 2(y-\overline{y})\cdot f'(h)
    
    
     Δw2​=η⋅∂w1​∂E​=η⋅2(y−y​)⋅f′(h)
    
    

    内存消耗的三方面

    输入数据

    很小,不做考量。

    256256的彩色图像:25625631 byte= 192KB

    模型参数

    较大,和模型复杂度有关。

    入门级的MNIST识别网络有6.6 million参数,使用32-bit浮点精度,占内存:6.6M 32 bit = 25MB 50层的ResNet有26 million参数,占内存:26M 32 bit = 99MB

    当然,你可以设计精简的网络来处理很复杂的问题。

    各层响应

    较大,同样和模型复杂度有关。

    50层的ResNet有16 million响应,占内存:16M*32bit = 64MB

    响应和模型参数的数量并没有直接关系。卷积层可以有很大尺寸的响应,但只有很少的参数;激活层甚至可以没有参数。

    – 这样看起来也不大啊?几百兆而已。
    – 往下看。

    batch的影响

    为了有效利用GPU的SIMD机制,要把数据以mini-batch的形式输入网络。
    如果要用32 bit的浮点数填满常见的1024 bit通路,需要32个样本同时计算。

    在使用mini-batch时,模型参数依然只保存一份,但各层响应需要按mini-batch大小翻倍。

    50层的ResNet,mini-batch=32,各层相应占内存:64MB*32 = 2GB

    卷积计算的影响

        H
    
    
        ×
    
    
        W
    
    
    
       H\times W
    
    
    H×W的输入图像为
    
    
    
    
        X
    
    
    
       X
    
    
    X,
    
    
    
    
        K
    
    
        ×
    
    
        K
    
    
    
       K\times K
    
    
    K×K的卷积核为
    
    
    
    
        R
    
    
    
       R
    
    
    R,符合我们直觉的卷积是这样计算的。
    
    

    对每一个输出位置,计算小块对位乘法结果之和。

          Y
    
    
          (
    
    
          h
    
    
          ,
    
    
          w
    
    
          )
    
    
          =
    
    
          ∑
    
    
    
    
            X
    
    
    
             k
    
    
             ,
    
    
             k
    
    
    
            s
    
    
    
           (
    
    
           h
    
    
           ,
    
    
           w
    
    
           )
    
    
           ⊙
    
    
           R
    
    
    
    
         Y(h,w) = \sum{X^s_{k,k}(h,w) \odot R}
    
    
      Y(h,w)=∑Xk,ks​(h,w)⊙R</p> 
    

    h = 1 : H , w = 1 : W h=1:H, w=1:W h=1:H,w=1:W
    其中, X k , k s ( h , w ) X^s_{k,k}(h,w) Xk,ks​(h,w)表示输入图像中,以 h , w h,w h,w为中心,尺寸为 K × K K\times K K×K的子图像。

    但是,这种零碎运算很慢

    在深度学习库中,一般会采用lowering的方式,把卷积计算转换成矩阵乘法

    首先,把输入图像分别平移不同距离,得到

          K
    
    
          2
    
    
    
    
        K^2
    
    
     K2个
    
    
    
    
         H
    
    
         ×
    
    
         W
    
    
    
        H\times W
    
    
     H×W的位移图像,串接成
    
    
    
    
         H
    
    
         ×
    
    
         W
    
    
         ×
    
    
    
          K
    
    
          2
    
    
    
    
        H\times W \times K^2
    
    
     H×W×K2的矩阵
    
    
    
    
    
          X
    
    
          ‾
    
    
    
    
        \overline{X}
    
    
     X。<br> 之后,把
    
    
    
    
         K
    
    
         ×
    
    
         K
    
    
    
        K\times K
    
    
     K×K的卷积核按照同样顺序拉伸成
    
    
    
    
    
          K
    
    
          2
    
    
    
         ×
    
    
         1
    
    
    
        K^2\times 1
    
    
     K2×1的矩阵
    
    
    
    
    
          R
    
    
          ‾
    
    
    
    
        \overline{R}
    
    
     R 
    
    
    
    
          Y
    
    
          =
    
    
    
           X
    
    
           ‾
    
    
    
          ⋅
    
    
    
           R
    
    
           ‾
    
    
    
    
         Y=\overline{X}\cdot \overline{R}
    
    
      Y=X⋅R</p> 
    
    
    

    输入输出为多通道时,方法类似,详情参见这篇博客

    在计算此类卷积时,前层响应

        X
    
    
    
       X
    
    
    X需要扩大
    
    
    
    
    
         K
    
    
         2
    
    
    
    
       K^2
    
    
    K2倍。
    
    

    50层的ResNet,考虑lowering效应时,各层响应占内存7.5GB

    使用低精度不能降内存

    为了有效利用SIMD,如果精度降低一倍,batch大小要扩大一倍。不能降低内存消耗。

    降内存的有效方法

    in-place运算

    不开辟新内存,直接重写原有响应。 复杂一些,通过分析整个网络图,可以找出只需要用一次的响应,它可以和后续响应共享内存。例如MxNet的memory sharing机制。

    综合运用这种方法,MIT在2016年的这篇论文能够把内存降低两到三倍。

    计算换存储

    找出那些容易计算的响应结果(例如激活函数层的输出)不与存储,在需要使用的时候临时计算。

    使用这种方法,MxNet的这个例子能够把50层的ResNet网络占用的内存减小四倍。

    类似地,DeepMind在2016年的这篇论文用RNN处理长度为1000的序列,内存占用降低20倍,计算量增加30%。

    百度语音在2016年的这篇论文同样针对RNN,内存占用降低16倍,可以训练100层网络。

    当然,还有Graphcore自家的IPU,也通过存储和计算的平衡来节约资源。

    Graphcore本身是一家机器学习芯片初创公司,行文中难免夹带私货,请明辨。

  • posted @ 2020-12-28 09:51  刘桓湚  阅读(2236)  评论(0编辑  收藏  举报