使用 Trace 实现 onnx 的导出 - 学习记录
使用 Trace 实现 onnx 的导出
- 一、使用 Trace 实现 onnx 的导出的流程
- 二、代码分解
-
- 2.1、定义模型
- 2.2、分析模型操作类型
- 2.3、构建钩子函数
-
- 2.3.1、定义 hook 函数
- 2.3.2、注册 Conv2d - hook 函数
- 2.3.3、注册 ReLU - hook 函数
- 2.3.4、注册 Add - hook 函数
- 三、完整导出 onnx 代码
前置学习内容：
Trace 功能实现
onnx的导出和分析 - onnx 注册自定义算子
一、使用 Trace 实现 onnx 的导出的流程
- 开始
- [定义模型]
- [创建钩子函数]
- [收集模型信息]
- [创建ONNX模型]
- [保存ONNX模型]
- [测试和验证]
- [结束]
二、代码分解
2.1、定义模型
class