Memory-efficient Implementation of DenseNets

改进densenet模型占用显存较大的问题。

对于特征的极致利用可以提高模型的表现能力,用时由于生成大量的intermediate feature(中间特征),因此存储这些intermediate feature会占用大量的显存。为了能够在GPU下跑更深的densenet网络,这篇文章通过对中间特征采用共享存储空间的方式降低了模型显存,使得在GPU显存限制下可以训练更深的densenet网络。当然这种共享部分存储也引入了额外的计算时间,因为在反向传播的时候需要重新计算一些层的输出,实际表现差不多增加了15%-20%的训练时间。

Memory-efficient Implementation of DenseNets

左图中,某个block的某一层卷积输入包含前面所有的卷积层的输出feature maps,这些feature map经过concate操作生成新的特征,然后经过normalization操作得到卷积的输入特征,这里concate和normalization操作生成的特征都重新开辟了存储空间。另外在torch中,反向传播也会开辟新的存储空间。在右图中,通过提前分配的shared memory storage和指针将这些intermediate feature(concate和bn操作生成的特征)存储在temporary storage buffers中,可大大减少存储量。

这种操作确实能够节省显存,不过也增加了计算时间。原因是:在前向计算的时候,concate得到的特征放在寄存器中,然后BN直接从寄存器中读取特征作为输入。在反向传播中要用到concate层输出或者bn层输出时,就要重新计算这两层的输出,因为寄存器的内容在前向传播过程中会被不断覆盖,这就增加了计算时间。如果是传统的方法,每个concate操作和bn操作生成的feature map都开辟了新的存储空间,那么反向传播的时候就可以直接读取在前向过程中保存的特征, 不需要重新生成了。

Pytorch是共享存储机制的

这种共享存储的做法的显存占用量并不会随着网络深度的增加而线性增加,主要是因为网络参数的存储大小要远远小于feature map的存储大小。

在Imagenet上原来只能训练161层的DenseNet,现在可以训练264层,从而有更深的效果。

在GPU显存限制下尽可能训练更深的densenet网络,通过共享intermediate feature的存储空间减少了显存占用,一定程度上增加了计算时间。

 

相关文章: