TNN框架解析

近一年多的时间&＃xff0c;一直都在做推理框架相关的工作&＃xff0c;期间也学习了很多东西&＃xff0c;特别是以前比较欠缺的CUDA编程相关内容。这一年来收获很多&＃xff0c;对推理框架认识也越来越深刻。以前纯做算法的时候&＃xff0c;越到后面&＃xff0c;越觉得无趣&＃xff0c;比如你经常会发现你辛辛苦苦调参了一周&＃xff0c;却发现毫无效果&＃xff0c;或者发现虽然能work,但是却不知道为什么能work&＃xff0c;还有算法工程师很多时候大部分时间都用来处理数据了&＃xff0c;总是做这样的工作有的时候感觉就是在浪费生命&＃xff0c;对自己的成长和职业发展没有太大意义。换了一个方向之后&＃xff0c;感觉心里踏实很多&＃xff0c;也很喜欢现在的方向&＃xff0c;虽然还有很多东西需要学习&＃xff0c;其实做推理框架期间&＃xff0c;以前算法的经验是非常有用的, 因为有些客户的算法需要适配我们自己的推理框架&＃xff0c;就需要懂算法的人&＃xff0c;如果你对算法比较了解&＃xff0c;那么解决问题就会方便很多&＃xff0c;比如现在很多CV方向的客户&＃xff0c;经常用的算法包括人脸检测领域的RetinaFace,人脸识别领域的arcFace,目标检测领域的YOLO系列和SSD系列&＃xff0c;这些算法我都是比较熟悉的算法&＃xff0c;所以做适配的时候&＃xff0c;就容易很多&＃xff0c;跟客户交流的时候也更方便。如果你不懂算法&＃xff0c;适配的时候就困难很多。懂算法还能够帮你更好的理解一个推理框架&＃xff0c;遇到框架中出现的问题&＃xff0c;也能够更好的进行分析和定位。这一年重点调研了AMD开源的推理框架MIGraphX,还有就是最近一直在调研的腾讯开源的推理框架TNN,这篇文章主要是对TNN框架做个简单的解析。

TNN架构设计层次还是非常清晰的&＃xff0c;整体架构与MIGraphX类似&＃xff0c;都包含了模型转换、低精度优化、编译优化和计算引擎。但是实现方式差异很大&＃xff0c;TNN的很多设计借鉴了Caffe的设计思想&＃xff0c;比如层的设计&＃xff0c;而MIGraphX则借鉴了TVM的设计思想&＃xff0c;更偏向于深度学习编译器的设计理念。关于两者的区别下面会进一步阐述。

tnn中有几个重要的数据类型&＃xff0c;包括&＃xff1a;

1. TNN&＃xff1a;对外的接口&＃xff0c;主要用于模型解析和创建网络
2. Instance&＃xff1a;网络实例&＃xff0c;包含了实际的tnn网络
3. AbstractModelInterpreter&＃xff1a;模型解析器的接口&＃xff0c;所以模型解析器都需要继承该接口
4. AbstractNetwork: 实际的TNN网络接口&＃xff0c;所有tnn中的网络都需要继承该接口
5. BaseLayer&＃xff1a;所有TNN中层的接口&＃xff0c;所有tnn中的层都需要继承该接口
6. AbstractDevice&＃xff1a;由于TNN支持不同的后端(比如arm,x86,CUDA)&＃xff0c;AbstractDevice为所有后端实现的接口
7. BlobManager:负责内存管理

下面看一下他们之间的关系。

TNN的几个关键类型的类图如上图所示&＃xff0c;虚线表示依赖关系&＃xff0c;比如AbstractNetwork依赖BaseLayer&＃xff0c;BlobManager和AbstractDevice。其中TNN类作为TNN整个类图框架中的最顶层的类&＃xff0c;是对外提供的一组接口&＃xff0c;用户在使用TNN的时候&＃xff0c;就通过TNN类对外的接口来进行模型解析和创建网络,TNN的init方法创建一个模型解析器并进行模型解析&＃xff0c;CreateInst方法负责创建网络。TNN类的CreateInst方法实际通过Instance类型来创建网络的。Instance类创建的网络保存在成员变量network中。TNN中的网络由层组成&＃xff0c;这一点跟caffe是一致的&＃xff0c;同时包含了内存管理模块&＃xff0c;所以上述类图中AbstractNetwork依赖于BaseLayer&＃xff0c;BlobManager和AbstractDevice&＃xff0c;下面重点分析一下network的这几个部分。

设备

TNN中的设备表示支持的各种后端&＃xff0c;包括arm,x86,cuda等。device的重要作用就是提供了不同后端的算子实现&＃xff0c;算子的具体某个后端的实现通过AbstractDevice的CreateLayerAcc()方法来创建。

层

TNN中层的概念与Caffe中的层的概念基本一致&＃xff0c;同样包含了caffe中层的reshape,forward等方法&＃xff0c;如果你对caffe很了解&＃xff0c;那么对TNN的层就很容易理解。层中的成员变量layer_acc为具体设备的实现&＃xff0c;该成员变量通过device的CreateLayerAcc()方法赋值实现不同后端的算子计算。注意&＃xff0c;TNN中cuda后端是通过tensorrt来实现的&＃xff0c;层的实现也是通过tensorrt的插件来实现的。

内存管理

推理框架中一个很重要的模块就是内存管理。TNN中通过BlobManager来实现内存管理。BlobManager又是通过内存池BlobMemoryPool来实现实际的内存管理。上述类图中&＃xff0c;Blob与caffe中的blob概念类似&＃xff0c;Blob中包含了blob的属性描述和内存数据&＃xff0c;BlobMemory保存了Blob的实际的内存&＃xff0c;并通过内存池来分配和释放。TNN中将Blob和BlobMemory进行解耦是为了更好的优化内存。

TNN中一个重要的设计模式&＃xff1a;工厂方法模式

TNN中大部分类型通过注册器注册到系统中&＃xff0c;而TNN中的注册用到了一个关键的设计模式&＃xff1a;工厂方法模式。

BaseLayer作为产品类的基类&＃xff0c;下面派生出各种产品&＃xff0c;包括卷积层&＃xff0c;Pooling层等。LayerCreator作为工厂类的基类&＃xff0c;包含了一个工厂方法CreateLayer(),注意&＃xff0c;这个方法返回的是BaseLayer*。这样才能够实现创建各种不同的产品。最后会通过一个注册器来实现自动注册。注册器的实现如下&＃xff1a;

template <typename T>
class TypeLayerRegister
public:
explicit TypeLayerRegister(LayerType type)
GetGlobalLayerCreatorMap()[type] &＃61; shared_ptr<T>(new T(type));

;


注册器在构造函数中实现产品和工厂的注册&＃xff0c;在构造函数中&＃xff0c;首先获取一个std::map,然后通过参数type和模板实现赋值&＃xff0c;从而实现注册的功能。最后在每个产品类中&＃xff0c;创建一个注册器的全局变量&＃xff0c;比如在卷积层的.cpp文件中创建一个全局变量&＃xff1a;

TypeLayerRegister<TypeLayerCreator<ConvLayer>> g_conv_register(LAYER_CONVOLUTION);


因为全局变量在程序启动的时候就会自动创建&＃xff0c;而创建的时候会调用构造函数&＃xff0c;从而实现自动注册的功能。

下面看一下常用的深度学习框架在op粒度,内存和动态shape方面的对比。

下面以TNN和MIGraphX为例&＃xff0c;详细讨论一下各方面的对比。

OP粒度

TNN在整体的设计方面参考了caffe的设计。下表列出了两者在一些数据类型上的对比&＃xff1a;

可以看出&＃xff0c;TNN中主要的数据类型都参考了caffe的设计&＃xff0c;在整体架构设计上也有很多地方参考了caffe。caffe属于第一代深度学习框架&＃xff0c;op的粒度比较粗&＃xff0c;所以TNN的op粒度也比较粗&＃xff0c;而MIGraphX的op粒度相比TNN来说要更细。下表列举了一些算子在两种框架中的实现差异&＃xff1a;

从上表可以看出&＃xff0c;TNN中通常可以使用一个层实现的计算&＃xff0c;在MIGraphX中可能需要好几个算子来实现。细粒度算子虽然比较灵活&＃xff0c;但是由于会增加重复访问内存的次数&＃xff0c;所以速度慢&＃xff0c;而粗粒度虽然灵活度不够&＃xff0c;但是速度快。这一点是TNN和MIGraphX的一个非常重要的区别。对于细粒度速度较慢的问题&＃xff0c;可以通过计算图优化提高速度&＃xff0c;典型的计算图优化包括算子融合&＃xff0c;算子等价&＃xff0c;删除公共子表达式。下表列出了MIGraphX的一些计算图优化。

上面提到了MIGraphX的设计参考了TVM&＃xff0c;MIGraphX中很多概念都借鉴了TVM。

1. 两者都通过pass这种类型来实现对计算图的优化
2. 两者都实现了像fuse_op,dead_code_elimination,eliminate_common_subexpression这些计算图的优化
3. 两者都通过lowering这个pass完成到更低级代码的转换

内存优化

内存的优化是推理框架的一个重要组成部分&＃xff0c;特别是对于内存较小的一些场景&＃xff0c;比如移动端设备。TNN和MIGraphX都对内存做了优化。下面我们看一下几个经典模型在TNN中和MIGraphX中的内存占用(注&＃xff1a;不包括算子参数占用的内存,只包含每层输出的内存)。

注&＃xff1a;测试模型的batchsize为1&＃xff0c;且使用FP32模式

从上表可以看出&＃xff0c;TNN优化后的内存占用平均为优化前的32%&＃xff0c;而MIGraphX优化后的内存占用平均为优化前的26%。由于两个框架对计算图优化存在较大差异&＃xff0c;所以不能仅从上述表格中的数据说明哪个内存优化效率更高。下面分析一下两个框架内存优化原理。

TNN的内存优化

TNN官方文档中对内存优化的介绍说&＃xff1a; 通过 DAG 网络计算图分析&＃xff0c; 实现无计算依赖的节点间复用内存&＃xff0c; 降低 90% 内存资源消耗。阅读完TNN内存优化部分的源码之后&＃xff0c;发现TNN对内存优化的方法其实很简单&＃xff0c;下面用一个非常简单的网络阐述一下TNN的优化过程&＃xff1a;

最左边是原始的网络结构&＃xff0c;右边为TNN的内存优化过程&＃xff0c;其中不同颜色代表了不同的内存块&＃xff0c;同一种颜色表示相同的一个内存块。在分析的过程中不考虑ReLU层。

TNN在优化内存的时候&＃xff0c;主要的思想就是&＃xff1a;创建一个候选内存块链表&＃xff0c;保存之前已经使用过的内存块&＃xff0c;然后在每次创建新层的时候&＃xff1a;为该层的输出分配内存的时候&＃xff0c;首先从候选内存块中找是否有合适的内存块&＃xff0c;如果找到了就直接用该内存块&＃xff0c;否则申请一块新的内存块&＃xff0c;然后查看该层的输入&＃xff0c;如果该层的输入不再被使用&＃xff0c;则放入候选内存块链表中&＃xff0c;供下面的新层使用。

下面分析一下上诉网络的内存优化过程&＃xff1a;

1. 首先创建一个候选内存块链表&＃xff0c;此时链表为空
2. 首先创建数据层&＃xff1a;为数据层创建一个内存块&＃xff0c;注意&＃xff1a;数据层的内存块是不能放入候选内存块链表中的&＃xff0c;因为这个内存块中保留着数据层中的数据&＃xff0c;在整个推理过程中都可能用到&＃xff0c;所以不能被覆盖
3. 创建conv1&＃xff1a;首先需要为该层输出分配内存&＃xff0c;先查找候选内存块是否有合适的可用的内存块&＃xff0c;发现此时没有&＃xff0c;那么就申请一个新的内存块&＃xff0c;我们用蓝色表示&＃xff0c;然后查看该层的输入是否可以被放入候选内存块链表中&＃xff0c;由于该层的输入是数据层&＃xff0c;所以不能&＃xff0c;此时候选内存块链表还是为空
4. 创建pool1&＃xff1a;首先需要为该层输出分配内存&＃xff0c;先查找候选内存块是否有合适的可用的内存块&＃xff0c;发现此时也没有&＃xff0c;那么就申请一个新的内存块&＃xff0c;我们用绿色表示&＃xff0c;然后查看该层的输入是否可以被放入候选内存块链表中&＃xff0c;由于该层的输入是conv1的输出&＃xff0c;而此时conv1的输出已不再被其他层使用&＃xff0c;所以可以放入候选内存块链表中
5. 创建conv2:首先需要为该层输出分配内存&＃xff0c;先查找候选内存块是否有合适的可用的内存块&＃xff0c;发现此时有蓝色内存块可以使用同时将蓝色内存块从链表中删除&＃xff0c;然后查看该层的输入是否可以被放入候选内存块链表中&＃xff0c;由于该层的输入是pool1的输出&＃xff0c;而此时pool1的输出已不再被其他层使用&＃xff0c;所以可以放入候选内存块链表中
6. 创建conv3:首先需要为该层输出分配内存&＃xff0c;先查找候选内存块是否有合适的可用的内存块&＃xff0c;发现此时有绿色内存块可以使用同时将绿色内存块从链表中删除&＃xff0c;然后查看该层的输入是否可以被放入候选内存块链表中&＃xff0c;由于该层的输入是conv2的输出&＃xff0c;而此时conv2的输出已不再被其他层使用&＃xff0c;所以可以放入候选内存块链表中
7. 创建conv4-2:首先需要为该层输出分配内存&＃xff0c;先查找候选内存块是否有合适的可用的内存块&＃xff0c;发现此时有蓝色内存块可以使用同时将蓝色内存块从链表中删除&＃xff0c;然后查看该层的输入是否可以被放入候选内存块链表中&＃xff0c;由于该层的输入是conv3的输出&＃xff0c;而此时conv3的输出同时被conv4-1层使用&＃xff0c;所以不能放入链表中&＃xff0c;所以此时链表为空
8. 创建conv4-1:首先需要为该层输出分配内存&＃xff0c;先查找候选内存块是否有合适的可用的内存块&＃xff0c;发现此时没有&＃xff0c;那么就申请一个新的内存块&＃xff0c;我们用表示红色表示&＃xff0c;然后查看该层的输入是否可以被放入候选内存块链表中&＃xff0c;由于该层的输入是conv3的输出&＃xff0c;而此时conv3的输出已经不被其他层使用了&＃xff0c;所以可以放入链表中&＃xff0c;此时链表中有绿色内存块
9. 创建prob1:首先需要为该层输出分配内存&＃xff0c;先查找候选内存块是否有合适的可用的内存块&＃xff0c;发现此时有绿色内存块可以使用同时将绿色内存块从链表中删除&＃xff0c;然后查看该层的输入是否可以被放入候选内存块链表中&＃xff0c;由于该层的输入是conv4-1的输出&＃xff0c;而conv4-1的输出已不再被其他层使用&＃xff0c;所以可以放入候选内存块链&＃xff0c;此时链表中有红色内存块

通过上述内存优化&＃xff0c;就完成了整个网络的内存分配工作&＃xff0c;可以看到一共使用了4个内存块&＃xff0c;显著提高了内存使用效率。

MIGraphX的内存优化

migraphx的内存优化通过图论中的图着色算法实现的。图着色的基本思想就是&＃xff1a;给定无向连通图G和m种不同的颜色。用这些颜色为图G的各顶点着色&＃xff0c;每个顶点着一种颜色。是否有一种着色法使G中每条边的2个顶点着不同颜色。若一个图最少需要m种颜色才能使图中每条边连接的2个顶点着不同颜色&＃xff0c;则称这个数m为该图的色数。

比如上图的色数为3,也就是说至少需要3中颜色就可以使得每条边的2个顶点颜色不同。联系上面的TNN内存优化策略&＃xff0c;我们可以发现&＃xff0c;TNN的内存优化策略其实就是让图中每条边的2个顶点颜色不同。上述TNN的示例中&＃xff0c;不算数据层&＃xff0c;则那张图的色度为3。虽然migraphx使用的内存优化算法实现和TNN不同&＃xff0c;但是核心思想是一样的&＃xff0c;就是让每条边的两个顶点颜色不同。从而达到内存复用的效果。

动态输入的支持

实际上完全的动态输入对内存是非常不友好的&＃xff0c;对于训练框架来说&＃xff0c;这个影响可能不大&＃xff0c;但是对于推理框架来说这个影响是非常大的&＃xff0c;比如caffe,可以在每次Forward的时候&＃xff0c;进行Reshape,重新进行内存分配和释放&＃xff0c;但是如果推理框架每次Forward的时候&＃xff0c;都需要进行内存分配和释放&＃xff0c;那么对推理效率的影响是非常大的。上面提到的TNN和MIGraphX的内存优化都是在有确定的输入大小的情况下才能够实现的。但是TNN是可以实现动态输入的&＃xff0c;但是MIGraphX目前是不支持的。下面我们看一下TNN是如何实现的。

TNN中创建网络的时候可以使用下面的一个接口&＃xff1a;

std::shared_ptr<Instance> CreateInst(NetworkConfig& config, Status& status,InputShapesMap min_inputs_shape, InputShapesMap max_inputs_shape);


注意&＃xff0c;这里有个max_inputs_shape参数&＃xff0c;这个参数的意思就是最大输入大小&＃xff0c;如果你需要支持动态输入&＃xff0c;则需要指定动态输入大小的最大尺寸是多少&＃xff0c;否则无法使用动态输入。一般情况下&＃xff0c;当网络输入大小改变了的时候&＃xff0c;对网络进行Forward的时候&＃xff0c;如果Blob的元素大小大于原来的大小&＃xff0c;则需要进行重新内存分配。但是TNN可以做到在网络输入大小改变之后&＃xff0c;不需要重新内存分配。这其中有个关键的地方&＃xff1a;创建网络的时候&＃xff0c;使用max_inputs_shape进行内存分配和优化。当使用max_inputs_shape进行内存分配和优化的时候&＃xff0c;网络中每个层分配的内存都是最大的&＃xff0c;这就保证了当输入为动态的时候&＃xff0c;每个层需要的内存是够用的&＃xff0c;不需要再重新分配了&＃xff0c;这样就可以大大减少内存分配的耗时&＃xff0c;上面TNN的内存优化示例中网络一共需要使用4个内存块&＃xff0c;这4个内存块都是使用max_inputs_shape计算出来的。TNN中当输入大小改变之后&＃xff0c;只需要通过每层的Reshape方法重新计算该层输出Blob的大小即可&＃xff0c;不需要重新内存分配了。TNN能够实现该机制还有一个重要原因就是将Blob的属性(BlobDesc)和数据(BlobMemory)进行了解耦&＃xff0c;使用内存池管理实际的数据。

虽然接触推理框架有一年的时间了&＃xff0c;但是觉得还是有很多地方需要学习。这篇文章只是对TNN的整体结构做了一个简单的介绍&＃xff0c;文中有什么描述不对的地方&＃xff0c;欢迎大家批评指正&＃xff0c;也欢迎做推理框架的朋友留言讨论。

2021-9-13 15:51:23

非常感谢您的阅读&＃xff0c;如果您觉得这篇文章对您有帮助&＃xff0c;欢迎扫码进行赞赏。