内容介绍:解决费显存的思路其实并不难,我们只需要预先分配一块缓存,供网络中所有的拼接层(Concatenation Layer)共享使用,这样 DenseNet 对内存的消耗便从平方级别降到了线性级别。在梯度反传过程中,我们再把相应卷积层的输出复制到该缓存,就可以重构每一层的输入特征,进而计算梯度。当然网络中由于 Batch Normalization 层的存在,实现起来还有一些需要注意的细节。为此我们专门写了一个技术报告(Memory-Efficient Implementation of DenseNets, https://arxiv.org/pdf/1707.06990.pdf)介绍如何提升 DenseNet 对内存的使用效率,同时还提供了 Torch, PyTorch, MxNet 以及 Caffe 的实现,代码参见:
Torch implementation: https://github.com/liuzhuang13/DenseNet/tree/master/models
PyTorch implementation: https://github.com/gpleiss/efficient_densenet_pytorch
MxNet implementation: https://github.com/taineleau/efficient_densenet_mxnet
Caffe implementation: https://github.com/Tongcheng/DN_CaffeScript
原DenseNet代码:https://github.com/pytorch/vision/blob/master/torchvision/models/densenet.py