Paddle2ONNX 架构设计
一、目标
1.1 背景
- AI工具库生态的碎片化：随着AI技术的快速发展，市场上涌现出了多种深度学习框架，如TensorFlow、PyTorch、PaddlePaddle等。每种框架都有其独特的优势和生态系统，但这也导致了AI工具库生态的碎片化。不同框架之间的模型和数据格式互不兼容，使得模型迁移和部署变得复杂和耗时。
- 跨框架和硬件平台的需求：在实际应用中，研究人员和工程师经常需要在不同的深度学习框架和硬件平台之间迁移模型。例如，某个模型可能在PyTorch中训练，但需要在TensorFlow或PaddlePaddle中进行推理；或者需要在不同的硬件平台（如GPU、CPU、边缘设备等）上部署模型。这种跨框架和硬件平台的需求推动了模型转换工具的研发。
- ONNX标准的兴起：ONNX（Open Neural Network Exchange）作为一种针对机器学习设计的开放式文件格式，旨在解决不同框架和硬件平台之间的互操作性问题。ONNX定义了一组与环境、平台均无关的标准格式，用于存储训练好的模型，包括模型的权重、结构信息以及每一层的输入输出等。这使得不同框架训练的模型可以统一转换为ONNX格式进行存储和交互。
1.2 价值
- 提高模型迁移和部署的效率：paddle2onnx工具能够将PaddlePaddle模型转换为ONNX格式，从而方便地将模型迁移到其他支持ONNX的深度学习框架或硬件平台上。这大大简化了模型迁移和部署的过程，提高了效率。
- **增强模型的可交互性：通过支持ONNX标准，paddle2onnx工具增强了PaddlePaddle模型与其他AI框架和硬件平台之间的可交互性。研究人员和工程师可以更容易地在不同环境中共享和部署模型。
- 促进AI生态的开放和共享：ONNX标准的广泛支持和应用促进了AI生态的开放和共享。paddle2onnx工具作为连接PaddlePaddle和ONNX的桥梁，有助于推动AI技术的普及和发展。
- 支持更多硬件平台：除了深度学习框架外，许多推理引擎或国产硬件也支持加载ONNX模型进行推理。通过paddle2onnx工具将PaddlePaddle模型转换为ONNX格式后，可以在这些硬件平台上进行部署和应用，进一步拓宽了模型的应用场景。
二、架构
2.1 概览
从架构实现上，paddle2onnx 大致可以分为三层：
- Python 层：定义了与用户息息相关的模型转换接口（如export、dygraph2onnx）；以及二次开发相关的API（如optimize，转fp16）等；
- Pybind 层：可以看作是中间映射层。paddle2onnx的核心功能都是在C++后端实现的，通过Pybind统一暴露给前端来转发调用；
- C++ 后端：主体功能的核心实现，自底而上依次包括Program的proto定义、模型解析Parser组件、算子转换规则定义、ONNX转换辅助组件、模型转换组件等；
2.2 Python 前端
此部分比较简单，主要包含如下几个Python文件：
command.py：通过argparse来定义paddle2onnx的命令行方式的各个功能参数，如model_dir、opset_version、export_fp16_model等convert.py：定义了export、dygraph2onnx的接口；前者是直接转发调用 Pybind 的 Export 函数，后者是先通过paddle.jit.save动转静，然后再调用export接口optimize.py：通过argparse来定义 optimize 接口的命令行参数，然后转发调用 Pybind 层的 optimize 函数convert_to_fp16.py：通过argparse来定义 convert_to_fp16 接口的命令行参数，然后转发调用 Pybind 层的 convert_to_fp16 函数
2.3 Pybind 中间层
由于 paddle2onnx 的核心功能都是在C++端实现的，故需要通过 PYBIND11_MODULE 来统一暴露必要的接口函数，具体实现是在 paddle2onnx/cpp2py_export.cc 里：
PYBIND11_MODULE(paddle2onnx_cpp2py_export, m) {
m.doc() = "Paddle2ONNX: export PaddlePaddle to ONNX";
m.def("export", [](....){});
m.def("optimize", [](....){
ONNX_NAMESPACE::optimization::OptimizePaddle2ONNX(
model_path, optimized_model_path, shape_infos);
});
m.def("convert_to_fp16", [](....){
ONNX_NAMESPACE::optimization::Paddle2ONNXFP32ToFP16(fp32_model_path,
fp16_model_path);
});
这里一并介绍下 paddle2onnx 的编译方式，可以使用pip install -e .的方式来编译安装；如果涉及到新算子规则的添加，可以参考 《Paddle2ONNX 开发指南》 。
如下对 CMakeLists.txt 里重要的信息进行阐述：
CMAKE_MINIMUM_REQUIRED(VERSION 3.16)
PROJECT(paddle2onnx C CXX)
set(CMAKE_CXX_STANDARD 17)
# Set max opset version for onnx if you build from other version of onnx this should be modified.
# Refer to https://github.com/onnx/onnx/blob/main/docs/Versioning.md#released-versions
add_definitions(-DMAX_ONNX_OPSET_VERSION=19)
# 递归匹配和收集所有.cc 文件，适当过滤
file(GLOB_RECURSE ALL_SRCS ${PROJECT_SOURCE_DIR}/paddle2onnx/*.cc ${PROJECT_SOURCE_DIR}/third_party/optimizer/onnxoptimizer/*.cc)
list(REMOVE_ITEM ALL_SRCS ${PROJECT_SOURCE_DIR}/paddle2onnx/cpp2py_export.cc)
list(REMOVE_ITEM ALL_SRCS ${PROJECT_SOURCE_DIR}/third_party/optimizer/onnxoptimizer/cpp2py_export.cc)
# 关键头文件
install(
FILES
${PROJECT_SOURCE_DIR}/paddle2onnx/converter.h
${PROJECT_SOURCE_DIR}/paddle2onnx/mappers_registry.h
DESTINATION include/paddle2onnx
)
if (BUILD_PADDLE2ONNX_PYTHON)
# 编译的对象
add_library(paddle2onnx_cpp2py_export MODULE ${PROJECT_SOURCE_DIR}/paddle2onnx/cpp2py_export.cc ${ALL_SRCS})
target_link_libraries(paddle2onnx_cpp2py_export PRIVATE -Wl,-force_load,$<TARGET_FILE:onnx>)
endif
2.4 C++ 后端
2.4.1 Proto 定义
Paddle2ONNX 主要负责将 Paddle 的 Inference 模型转换为 ONNX 格式，便于开发者将 Paddle 模型扩展到支持 ONNX 部署的框架上进行推理。在飞桨3.0 之前，*.pdmodel 文件是通过Protobuf定义的（详见framework.proto），为了能够与主框架部分解耦，paddle2onnx 仓库里也引用了一份proto文件（详见 p2o_paddle.proto）
关于 paddle2onnx 里的 proto 文件相对主框架滞后更新常会引发一种报错：Paddle2ONNX/protobuf/src/google/protobuf/message_lite.cc:133] Can't parse message of type "paddle2onnx.framework.proto.ProgramDesc" because it is missing required fields xxx
paddle2onnx 直接依赖了onnx组件，故转换的目标结果是 ONNX_NAMESPACE::ModelProto 对象，其内在包含了 NodeProto 、 ValueInfoProto 等节点结构。
2.4.2 PaddleParser
paddle2onnx 接收的是飞桨导出的模型文件，一般包括 *.pdmodel 模型文件和 *.pdparams 参数文件。因此首先需要将磁盘文件反序列化为 C++ ProgramDesc 对象，主要过程包括：
LoadProgram：构建空的paddle2onnx::framework::proto::ProgramDesc对象，读取 pdmodel 二进字节流，使用ParseFromString反序列化LoadParams：由于飞桨的参数保存过程是遵循特定的格式，比如对于每个参数，按照类似 name、dtype、buffer_size、buffer_data等顺序存储，解析也是按照这个顺序；实现上是与Paddle主框架中的save_combine算子Kernel一致。在2onnx中会将每个参数构造为Weight对象，内含buffer data。InitBlock：主要包括3个子步骤：GetBlocksVarName2Id：解析所有 Block 中的 VarDesc 的name信息GetBlockOps：解析所有Block中所有的 OpDesc 信息GetGlobalBlockInputOutputInfo：解析Block[0]中的全局输入输出VarDesc，转为TensorInfo（通过feed、fetch 算子来识别）
2.4.3 ModelExporter
ModelExporter 提供了Run函数接口实现将飞桨的模型转换为ONNX形式，具体的职责是一次交给 PaddleParser 和 OnnxHelper 来做的。前者负责解析Desc信息，后者负责逐个调用算子对应的转换规则，实现到ONNX ModelProto的转换。
我们主要关注几个逻辑上的要点：
- 算子支持度：在对算子转换之前，会先调用
CheckIfOpSupported扫描Program中的所有算子，判断是否都支持。 - 参数的处理：对于所有的 Parameter，转为ONNX之后都会关联一个
ConstantOp - 算子的处理：每个算子都会在
Mapper规则实现里转为1个或多个NodeProto，具体取决于飞桨的算子定义与ONNX的差异性
2.4.4 算子转换规则 Mapper
算子转换规则是 paddle2onnx 的最核心的部分，其代码量也是最多的，主要定义在 paddle2onnx/mapper 目录中，并按照领域进行了子目录的划分。
此处我们以 Dropout 算子为例，如果要添加此算子的支持，我们需要写一个 dropout.h 头文件，可以看出：
- 需要统一继承
Mapper基类 - 实现
GetMinOpset()函数，来定义支持的opset_version - 按照
opset_version实现对应版本的OpsetX()函数 - 派生子类可以有自己的
Attribute成员
#pragma once
#include <vector>
#include "paddle2onnx/mapper/mapper.h"
namespace paddle2onnx {
class DropoutMapper : public Mapper {
public:
DropoutMapper(const PaddleParser& p, OnnxHelper* helper, int64_t block_id,
int64_t op_id)
: Mapper(p, helper, block_id, op_id) {
GetAttr("dropout_implementation", &dropout_implementation_);
}
int32_t GetMinOpset(bool verbose = false);
void Opset7();
private:
float dropout_prob_ = 0.0;
std::string dropout_implementation_ = "upscale_in_train";
};
然后需要在 dropout.cc 实现上面两个函数：
- 从
GetMinOpset可以看出：并非飞桨算子所有的功能都能转换为ONNX，根源是两个框架下算子的功能集合可能有差别 - 在这个例子中
GetMinOpset返回了7，故需要实现Opset7函数，最终是总会调用OnnxHelper->MakeNode()
int32_t DropoutMapper::GetMinOpset(bool verbose) {
if (dropout_implementation_ != "downgrade_in_infer" &&
dropout_implementation_ != "upscale_in_train") {
Error() << "Drop out type: " << dropout_implementation_
<< " is not supported yet." << std::endl;
return -1;
}
if (dropout_implementation_ == "downgrade_in_infer") {
if (IsAttrVar("dropout_prob") &&
!IsConstant(GetAttrVar("dropout_prob")[0])) {
Error() << "While Attribute(dropout_prob)'s type is Tensor, it's not "
"supported "
"unless it's a constant tensor when dropout_implementation is "
"downgrade_in_infer."
<< std::endl;
return -1;
}
}
return 7;
}
void DropoutMapper::Opset7() {
auto input_info = GetInput("X");
auto output_info = GetOutput("Out");
if (dropout_implementation_ == "upscale_in_train") {
helper_->MakeNode("Identity", {input_info[0].name}, {output_info[0].name});
} else {
if (IsAttrVar("dropout_prob")) {
auto prob_info = GetAttrVar("dropout_prob");
std::vector<float> temp;
TryGetValue(prob_info[0], &temp);
dropout_prob_ = temp[0];
} else {
GetAttr("dropout_prob", &dropout_prob_);
}
std::string scale_node = helper_->Constant(
{}, GetOnnxDtype(input_info[0].dtype), 1 - dropout_prob_);
helper_->MakeNode("Mul", {input_info[0].name, scale_node},
{output_info[0].name});
}
}
最后一步，在 dropout.cc 中调用 REGISTER_MAPPER(dropout, DropoutMapper) 来注册。
三、关键问题
3.1 不同算子版本管理机制
由于所有的算子规则实现都会继承 Mapper 基类，我们看下基类里的主体逻辑是怎么管理不同 opset_version 的ONNX 算子的：
- Opset Version 需要
>= 7 - Run 本质是一个 Dispatcher 的设计，根据Opset Version来分发
- 每个 Version 的默认实现是调用下一个版本，即 Fallback 逻辑
class Mapper {
public:
// the return value in [7, MAX_ONNX_OPSET_VERSION], represent the minimum
// opset_version
// if return value < 0, means the op is not supported.
virtual int32_t GetMinOpset(bool verbose) { return 7; }
void Run() {
int32_t opset_version = helper_->GetOpsetVersion();
Assert(opset_version >= 7 && opset_version <= MAX_ONNX_OPSET_VERSION,
"[Paddle2ONNX] Only support opset_version in range of [7, " +
std::to_string(MAX_ONNX_OPSET_VERSION) + "].");
if (opset_version == 19) {
Opset19();
} else if (opset_version == 18) {
Opset18();
// ..... 省略
}else {
Opset7();
}
}
virtual void Opset19() { Opset18(); }
virtual void Opset18() { Opset17(); }
// ..... 省略
virtual void Opset7() {
// Raise Error
}
};
在具体的Mapper子类里，至少要实现一个 OpsetX 的函数逻辑，可以有多个，比如 SliceMapper 实现了 Opset7 和 Opset10 两种转换规则：
int32_t SliceMapper::GetMinOpset(bool verbose) {
if (HasInput("StartsTensorList") || HasInput("EndsTensorList") ||
HasInput("StridesTensorList")) {
Logger(verbose, 10)
<< "While has input StartsTensorList/EndsTensorListStridesTensorList, "
<< RequireOpset(10) << std::endl;
return 10;
}
if (HasInput("StartsTensor")) {
auto info = GetInput("StartsTensor");
if (!IsConstantInput("StartsTensor")) {
Logger(verbose, 10)
<< "While has input StartsTensor, and it's not a constant tensor, "
<< RequireOpset(10) << std::endl;
return 10;
}
}
if (HasInput("EndsTensor")) {
auto info = GetInput("EndsTensor");
if (!IsConstantInput("EndsTensor")) {
Logger(verbose, 10)
<< "While has input EndsTensor, and it's not a constant tensor, "
<< RequireOpset(10) << std::endl;
return 10;
}
}
if (HasInput("StridesTensor") || strides_.size() > 0) {
Logger(verbose, 10) << "While has strides, " << RequireOpset(10)
<< std::endl;
return 10;
}
return 7;
}
3.2 控制流支持情况
paddle2onnx 是有限地支持控制流算子，其转换规则是定义在 paddle2onnx/mapper/loop.cc 中，是否支持的逻辑是 IsLoopSupported
- 只支持部分场景的 While 算子。如不能包含
TensorArray，且输入输出的个数必须一样 - 完全不支持控制流 If 算子
bool ModelExporter::IsLoopSupported(const PaddleParser& parser,
const int64_t& block_id,
const int64_t& op_id) {
auto x_info = parser.GetOpInput(block_id, op_id, "X");
auto out_info = parser.GetOpOutput(block_id, op_id, "Out");
auto cond_info = parser.GetOpInput(block_id, op_id, "Condition");
std::set<std::string> input_names;
for (size_t i = 0; i < x_info.size(); ++i) {
input_names.insert(x_info[i].name);
}
input_names.insert(cond_info[0].name);
for (size_t i = 0; i < out_info.size(); ++i) {
auto iter = input_names.find(out_info[i].name);
if (iter == input_names.end()) {
P2OLogger() << "Cannot find output:" << out_info[i].name << " in input tensors while converting operator 'while', Paddle2ONNX doesn't support this situation now." << std::endl;
return false;
}
}
for (size_t i = 0; i < x_info.size(); ++i) {
if (x_info[i].is_tensor_array) {
P2OLogger() << "LodTensorArray is not supported." << std::endl;
return false;
}
}
return true;
}
