介绍ATC工具的功能架构以及模型转换过程中,各组件的交互流程。
昇腾张量编译器(Ascend Tensor Compiler,简称ATC)是异构计算架构CANN体系下的模型转换工具, 它可以将开源框架的网络模型以及Ascend IR定义的单算子描述文件(json格式)转换为昇腾AI处理器支持的.om格式离线模型。其功能架构如图1所示。
模型转换过程中,ATC会进行算子调度优化、权重数据重排、内存使用优化等具体操作,对原始的深度学习模型进行进一步的调优,从而满足部署场景下的高性能需求,使其能够高效执行在昇腾AI处理器上。
其中:
用户也可以将开源框架网络模型转换后的离线模型转成json文件,方便查看相关参数;也可以直接将开源框架网络模型通过ATC工具转成json文件,查看相关参数。
Ascend IR定义的单算子描述文件(json格式)通过ATC工具进行单算子编译后,转成适配昇腾AI处理器的单算子离线模型,然后上传到板端环境,通过AscendCL接口加载单算子模型文件用于验证单算子功能。
下面以开源框架网络模型转换为.om离线模型为例,详细介绍模型转换过程中与周边模块的交互流程。
根据网络模型中算子计算单元的不同,分为TBE(Tensor Boost Engine)算子、AI CPU算子,TBE算子在AI Core上运行,AI CPU算子在AI CPU上运行。在TBE算子、AI CPU算子的模型转换交互流程中,虽然都涉及图准备、图拆分、图优化、图编译等节点,但由于两者的计算单元不同,因此涉及交互的内部模块也有所不同,请参见下图。关于算子类型、基本概念等详细介绍请参见《TBE&AI CPU自定义算子开发指南》。
原图优化时:GE(Graph Engine,基于昇腾AI软件栈对不同的机器学习框架提供统一的IR接口,对接上层网络模型框架,例如Tensorflow、PyTorch等,GE的主要功能包括图准备、图拆分、图优化、图编译、图加载、图执行和图管理等。)向FE发送图优化请求,并将图下发给FE,FE匹配融合规则进行图融合,并进行算子选择,选择优先级最高的算子类型进行算子匹配,最后将优化后的整图返回给GE。
优化后的子图合并为整图,再进行整图优化。
另外,GE将整图下发给AI CPU Engine,AI CPU Engine读取算子信息库,匹配算子支持的format,并将format返回给GE。