GPU 显存 - Caffe 内存优化

Caffe - 显存优化

1. Caffe - memory multiloading 功能

原文 - Squeezing Memory out of Caffe

Caffe 的一个优化版本 - caffe-yjxiong. 测试了一下, ResNet101 相对于官方 caffe, 能够明显节省显存占用, batchsize 可以增加很多.

显存优化的 Caffe 主要实现的功能: memory multiloading

在深度网络训练和测试时,可以显著地节省内存.

训练时,节省一半内存; 测试时, 使用 95% 的内存.

该内存优化的核心思想比较简单:

如果有些消费者请求一些资源, 但他们并不需要同时请求所有的资源. 一种理想情况是, 只让他们在他们真正需要的时候再使用资源, 并在结束时返回资源. 在训练深度网络时, GPU显存资源是有限的. 资源消费者即是网络中的网络层(layers/operations), 在 GPU 保存着训练时的中间结果(intermediate results).

例如, 采用标准 BP 算法训练 CNNs 时, 最主要的一种内存消耗是, 保存中间激活值和其梯度(也叫作敏感度sensitivities)( intermediate activation values and the gradients values w.r.t. intermediate activation).

如果将一次 forward/backward 迭代看做一个周期(cycle), 那么一个重要发现是: 在整个周期中不是所有的值都需要保存的.

具体来说, 对于单个网络层, 在 feed-forwad 过程中, 当得到该网络层的输出值后, 则网络层的输入激活值就没用了. 在 BP 过程中, 当得到网络层参数的梯度和输入敏感度(input sensitivities) 后, 则网络层的输入激活值和敏感度sensitivities 就没用了. 也就是说, 每个内存块只有在计算时才占用物理内存； 且, 不再使用时, 就可以释放物理内存.

内存优化核心思想的实现可以共享存储资源并节省内存. 在 TensorFlow 或 Parrots 等深度框架中, 其主要是通过动态调度一个预先分配的内存池来实现的, 类似于 OS 对于物理 RAM 的调度(dynamically scheduling a pre-allocated memory pool as the OS does with the physical RAM). 但是, 在 Caffe 中, 如何实现呢? 由于 Caffe 具有完全不同的内存模型, 其每个内存块是由 OS/GPU 来自动分配的. 在内存池中实现类似功能是比较棘手的.

相反地, 采用了一种替代方法. 假设网络结构是固定的(大部分情况都是这样), 在运行第一次迭代前, 来判断一个内存块的生命周期(life-cycle). 这里称之为 “dry-run”, 即, 模拟计算流以寻找可以共享存储的 SyncedMem-s(Caffe 中的基本内存单元). 当找出 SyncedMem-s 后, 即可将其分配到单个内存插槽(single memory slot). 在实际训练和测试时, memory multiloading 的实现是将共享 SyncedMem-s 都由相同内存块来存储和读取的. 由于在 “dry-run” 中已经确定了依赖关系, multiloading 是不会损坏数据, 且能节省大量内存使用的.

最后, 在实现 Caffe 的 memory multiloading 功能时, 发现的一个有意思的问题. Caffe 有一些网络层不进行计算, 只共享网络层的输入和输出 blobs(ShareData/ShareDiff). 这种情况在某些时候可能破坏 “dry-run” 处理所获得的依赖关系. 因此, 需要先让这些网络层告诉框架网络层所做的事情( first make these layers tell the framework that they doing this.) 在某些情况, 这些网络层是堆叠在一起以构建特定的结构. 由于会引入递归分享(recursive sharing, 一系列的 blobs 共享其激活值/梯度内存块), 导致问题更加棘手. 最终, 采用的方案是, 著名的 并查集(或不相交集合)数据结构( “union-find” (disjoint-set) data structure) 来使 memory multiloading 的实现更加干净和优雅.

memory multiloading 功能实现的文件在:

https://github.com/yjxiong/caffe/blob/action_recog/src/caffe/net.cpp#L1088

2. 内存优化 Memory Optimization

2.1 使用方法

该版本 Caffe 中, 内存优化是由网络配置中的mem_param 来实现的.

例如:

name: "example_net"
mem_param {
  optimize_train: true
  optimize_test: true
  exclude_blob: "fc1"
  exclude_blob: "fc2"
}

• optimize_train - 是否在训练网络中进行优化. 如果设为 True, 可以节省训练时一半的内存.
• optimize_test - 是否在测试网络中进行优化. 如果设为 True, 可以节省测试时大部分内存.
• exclude_blob - 指定优化处理中所排除的 blobs. 排除一个 blob, 可以确保该 blob 可以被 Python/Matlab 接口安全的读取. 默认情况, 网络所有的 inputs, outputs 和 losses 都是被排除的.

2.2 工作原理

节省内存的原因是, 对于训练的每次 forward 和 backward, 前一 blobs 使用的内存会被后面 blobs 重用.

内存优化模块通过确定网络中 blobs 的依赖关系来工作的. 有时称之为 multiloading.

具体来说, 在网络训练开始前, 运行一次 “dry-run” 来确定重用 blob 内存块的方式. 这是一个静态优化过程.

另一方面, Parrots 深度学习框架是通过动态调度内存使用的, 具有更优的内存节省和更好的灵活性.