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 深度学习框架是通过动态调度内存使用的, 具有更优的内存节省和更好的灵活性.

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏MixLab科技+设计实验室

自己动手做一个识别手写数字的web应用02

继续上文。 自己动手做一个识别手写数字的web应用01 01 再次进入docker容器 接着上一篇文章,我们继续使用上次新建好的容器,可以终端输入 : d...

45470
来自专栏利炳根的专栏

学习笔记TF066 : TensorFlow 移动端应用,iOS、Android系统实践

移动端应用原理。移动端、嵌入式设备应用深度学习方式,一模型运行在云端服务器,向服务器发送请求,接收服务器响应;二在本地运行模型,PC训练模型,放到移动端预测。向...

95600
来自专栏瓜大三哥

图像分割(四)

图像分割(四) 之基于FPGA的局部自适应分割 子模块设计 窗口缓存模块win_buf 本模块不做任何算法上的处理,只是负责将当前输入像素的二维窗口元素缓存并组...

25780
来自专栏Python小屋

Python使用系统聚类方法进行数据分类案例一则

首先解释一下为啥最近发的文章中代码都是截图而不是文本,这样做主要是希望大家能对着代码敲一遍而不是直接复制运行得到结果就算了,这样可以加深印象,学到更多东西。当然...

47140
来自专栏MelonTeam专栏

Windows下Tensorflow的环境安装和Hello world

导语 深度学习 数字识别 Tensorflow 环境搭建 本人对深度学习是0基础,python也是没有用过,最近很流行深度学习,因此也想学习一下,...

27250
来自专栏磨磨谈

ceph 的crush算法 straw

很多年以前,Sage 在写CRUSH的原始算法的时候,写了不同的Bucket类型,可以选择不同的伪随机选择算法,大部分的模型是基于RJ Honicky写的RUS...

38620
来自专栏祝威廉

MLSQL如何支持部署SKLearn,Tensorflow,MLLib模型提供API预测服务

部署成API服务时,除了要把raw数据特征化成向量外,研发还要想着怎么加载模型,产生模型的框架五花八门,比如Tensorflow,SKlearn,Spark M...

16140
来自专栏深度学习那些事儿

pytorch中读取模型权重数据、保存数据方法总结

pytorch中保存数据策略在长时间的深度训练中有很大的作用,我们可以通过保存训练好的权重,然后等到下次使用的时候再取出来。另外我们也可以通过迁移学习使用别人训...

9.6K80
来自专栏木子昭的博客

教你精确调整gif图片尺寸

按照以上方法, 我们可以对任何gif的尺寸进行精确调整, 对于这类实用小技巧,可以点赞记录一下, 以后或许会用到~

12930
来自专栏王小雷

降低PNG图片存储大小方法、图片压缩方法

降低PNG图片存储大小方法,图片压缩方法,如何降低PNG图片存储大小?前提是分辨率和尺寸大小不变,图形的透明部分不变。请看如下办法,亲测可用。 1. 将PNG图...

39180

扫码关注云+社区

领取腾讯云代金券