目录
背景
随着越来越多的AI算法在端侧的部署(部分是常驻系统的),系统的内存资源消耗变得越来越严重,有些算法只需要几M的运行内存,有些则可能要几百M的运行内存,如果能够将AI算法的总运行内存有效降低,对于整个系统会是一个巨大的优化。在所有问题中,有两类亟需我们使用内存复用算法进行优化:
- 模型内存占用大:个别模型的输入为(3 * 1080 * 1080)* 4 / (1024 * 1024) = 13.34MB,且模型层数很多,如果不采用内存复用的话,运行内存通常要300M甚至更大
- 运行内存资源有限:大多数嵌入式设备可用的内存资源非常紧缺
内存优化方法简介
内存优化的思路有很多,主要是以下几种:
- 模型剪枝、蒸馏:去掉模型中的部分权重
- 模型量化:将fp32的数据量化成int16或者int8
- 内存复用:复用为Feature Map分配的内存
前两种方法都是独立于模型部署,第三种方法则需要在模型部署时执行,本文会重点介绍如何利用内存复用算法,优化模型的运行内存。
内存复用算法设计
注:
- 笔者写了一个转换工具,用于分析模型的结构、计算内存复用的策略,下面分析中的很多数据都是来源于该工具
- 本文定义了一个指标MRR(Memory Reduction Ratio),来衡量内存优化算法的效果,MRR表示一个模型实际分配的内存,相对于它不采用内存复用所分配的内存,降低的比率。
- feature_map_len,表示模型中每一层的输出尺寸,例如,Mobilenet V1的第一个节点的输出维度为3x224x224,那么它的feature_map_len就是3*224*224=150528,假如每个元素都是float32的,那么这个Feature Map占的内存就是150528 * 4 /(1024 * 1024)= 0.57M
- 以下分析采用的都是开源模型: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的序号