深度学习计算框架综述（六）内存复用算法设计

目录

背景

内存优化方法简介

内存复用算法设计

1 无分支模型

2 单输入、单输出、有分支模型

3 多输入或多输出模型

内存复用算法流程概述

讨论

---

背景

随着越来越多的AI算法在端侧的部署（部分是常驻系统的），系统的内存资源消耗变得越来越严重，有些算法只需要几M的运行内存，有些则可能要几百M的运行内存，如果能够将AI算法的总运行内存有效降低，对于整个系统会是一个巨大的优化。在所有问题中，有两类亟需我们使用内存复用算法进行优化：

1. 模型内存占用大：个别模型的输入为（3 * 1080 * 1080）* 4 / (1024 * 1024) = 13.34MB，且模型层数很多，如果不采用内存复用的话，运行内存通常要300M甚至更大
2. 运行内存资源有限：大多数嵌入式设备可用的内存资源非常紧缺

内存优化方法简介

内存优化的思路有很多，主要是以下几种：

1. 模型剪枝、蒸馏：去掉模型中的部分权重
2. 模型量化：将fp32的数据量化成int16或者int8
3. 内存复用：复用为Feature Map分配的内存

前两种方法都是独立于模型部署，第三种方法则需要在模型部署时执行，本文会重点介绍如何利用内存复用算法，优化模型的运行内存。

 

内存复用算法设计

注：

1. 笔者写了一个转换工具，用于分析模型的结构、计算内存复用的策略，下面分析中的很多数据都是来源于该工具
2. 本文定义了一个指标MRR(Memory Reduction Ratio)，来衡量内存优化算法的效果，MRR表示一个模型实际分配的内存，相对于它不采用内存复用所分配的内存，降低的比率。
3. feature_map_len，表示模型中每一层的输出尺寸，例如，Mobilenet V1的第一个节点的输出维度为3x224x224，那么它的feature_map_len就是3*224*224=150528，假如每个元素都是float32的，那么这个Feature Map占的内存就是150528 * 4  /（1024 * 1024）= 0.57M
4. 以下分析采用的都是开源模型：Mobilenet V1、Mobilenet V2、RetinaFace，都可以在Github上找到Pretrained Model

我们先来看常见的几类模型：

1 无分支模型

我们以Mobilenet V1为例，如下图所示（图中只展示了一部分，VCAP在进行模型转换的时候，会将Batchnorm、Relu合并到Conv中，同时使用有向无环图(DAG)拓扑排序算法给每一层打上ID，确定每一层的执行顺序）：

我们先来看一下，假如不采用内存复用，需要分配多大的运行内存。将模型作为输入传到我们的工具中，得到如下输出，如不采用内存复用，需要分配的运行内存约为20M：

但如果我们能够找到所有层中输出尺寸最大的一层，把这一层的长度记为max_feature_map_len，那么我们只需要分配两块大小为max_feature_map_len的内存就够了，我们设计了一个结构体MemBlock来表示内存块：

typedef struct mem_block {
  int block_id;
  int mem_size;
  std::vector<int> layer_ids;
  int max_size_in_layers;
  bool used = false;
} MemBlock;

其中各成员的含义如下：

• block_id，MemBlock的序号
• mem_sizeÿ