superpoint的onnx转换全流程

流程"/>

superpoint的onnx转换全流程

Tensorflow转ONNX - it610

转PB文件（Graph Freezing）

如上所述，常规TF的模型导出方法会将网络信息与权重信息分开存储在不同文件当中，这在部署时候不是很方便。官方提供了一种Freeze Graph的方式，用于将模型相关信息统统打包到一个*.pb文件当中。

官方提供了相关工具freeze_graph，一般安装完TensorFlow后会自动添加到用户PATH相应的bin目录下，如果没有找到的话可以去TensorFlow源码tensorflow/python/tools/free_graph.py这个位置去找一下，或者直接通过命令行导入module的方式调用。

举例如下，如果有多个输出节点，用逗号隔开：

# 1.直接调用
freeze_graph --input_graph=/home/mnist-tf/graph.proto \--input_checkpoint=/home/mnist-tf/ckpt/model.ckpt \--output_graph=/tmp/frozen_graph.pb \--output_node_names=fc2/add \--input_binary=True# 2. 通过调用module的方式
python -m tensorflow.python.tools.freeze_graph \--input_graph=my_checkpoint_dir/graphdef.pb \--input_binary=true \--output_node_names=output \--input_checkpoint=my_checkpoint_dir \--output_graph=tests/models/fc-layers/frozen.pb

其中有一个比较困扰的点在于需要准确的知道输出节点的节点名称，我的做法是通过tf.get_default_graph().as_graph_def().node来得到各节点信息，然后再从中查看具体的输出节点名。

print([tensor.name for tensor in tf.get_default_graph().as_graph_def().node])

onnx                  1.10.1
onnx-simplifier       0.3.6
onnxoptimizer         0.2.6
onnxruntime           1.8.1

完整的SSD-tf版本转换为onnx示例：
.ipynb

本文所示示例为superpoint

magic-point网络训练模拟数据input(h,w,c)=(120,160,1)
magic-point网络导出coco数据的GT，使用input(h,w,c)=(240,320,1)
super-point网络训练coco数据input(h,w,c)=(240,320,1)
super-point网络inference的时候使用input(h,w,c)=(480,640,1)

项目作者给出了网络三个阶段保存的模型
MagicPoint (synthetic) mp_synth-v11
MagicPoint (COCO) mp_synth-v11_ha1_trained
SuperPoint (COCO) sp_v6
我们通过导出onnx已经确定了他们的input和output
mp_synth-v11
input: image
output: pred  logits  prob  prob_nms
mp_synth-v11_ha1_trained
input: image
output: pred  logits  prob  prob_nms
sp_v6
input: image
output: descriptors  prob_nms  prob  descriptors_raw pred logits

参考：
1.模型checkpoint格式

1.5 如果你的模型是frozen_inference_graph.pb格式，就需要且不知道input output的name，
就比较麻烦，这里我们不考虑这种情况
.md

2.将checkpoint格式转换为saved model格式(pb)
.py

3.将saved model格式转换为onnx

记得选择适合自己网络的opset
可参考：
.4.0/docs/OperatorKernels.md
.md

python -m tf2onnx.convert --saved-model saved_models/model/ --output saved_models/model/saved_model.onnx --opset 11

4.得到onnx后发现，无法正常显示网络输入和网络输出

5.安装onnx
通过onnx打印网络所有的input和output

import onnxname = "X.onnx"
onnx_file = name
model = onnx.load(onnx_file)
for m in range(len(model.graph.input)):print(model.graph.input[m].name)print("\n")for m in range(len(model.graph.output)):print(model.graph.output[m].name)

6.下载onnx-simplifier

5.通过onnx-simplifier，指定input=(N,H,W,C) ps: data_format='channels_last'
导出simplifier后的onnx，可看到网络每层的shape

import onnx
from onnxsim import simplifyonnx_file = "mp_synth-v11.onnx"
sim_onnx_path = "mp_synth-v11__simplify.onnx"model = onnx.load(onnx_file)
onnx.checker.check_model(model)# print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))# 列表里的 shape 改成自己对应的
model_simp, check = simplify(model, input_shapes={"image":[1, 120, 160, 1]}) onnx.save(model_simp, sim_onnx_path)print("Simplify onnx done !")

参考：

其它一些参考

ONNX学习笔记 - 知乎

如何在Netron中看到每一层的shape_BokyLiu的博客-CSDN博客

【模型加速】PointPillars模型TensorRT加速实验(2)_昌山小屋的博客-CSDN博客

onnx2pytorch和onnx-simplifer新版介绍 - 知乎

torch.onnx — PyTorch 1.9.1 documentation

torch.onnx — PyTorch 1.9.1 documentation

GitHub - microsoft/onnxjs: ONNX.js: run ONNX models using JavaScript

GitHub - microsoft/onnxjs-demo: demos to show the capabilities of ONNX.js

superpoint导出pb

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "7"
import yaml
import argparse
import logging
from pathlib import Pathlogging.basicConfig(format='[%(asctime)s %(levelname)s] %(message)s',datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO)
import tensorflow as tf  # noqa: E402from superpoint.models import get_model  # noqa: E402if __name__ == '__main__':with open("config.yml", 'r') as f:config = yaml.load(f)config['model']['data_format'] = 'channels_last'def mkdir_os(path):if not os.path.exists(path):os.makedirs(path)export_root_dir = "onnx/"mkdir_os(export_root_dir)export_dir = "onnx/pb_model_onlytwo_NoneH640W480C1"checkpoint_path = "model"'''data_shapesuperpoint/models/base_model.py# Prediction network with feed_dictif self.data_shape is None:self.data_shape = {i: spec['shape'] for i, spec in self.input_spec.items()}self.pred_in = {i: tf.placeholder(spec['type'], shape=self.data_shape[i], name=i)for i, spec in self.input_spec.items()}self._pred_graph(self.pred_in)'''with get_model(config['model']['name'])(data_shape={'image': [None, 640, 480, 1]},**config['model']) as net:net.load(str(checkpoint_path))# tf.saved_model.simple_save(#         net.sess,#         str(export_dir),#         inputs=net.pred_in,#         outputs=net.pred_out)# name:logits# name:prob# name:descriptors_raw# name:descriptors# name:prob_nms# name:predtf.saved_model.simple_save(net.sess,str(export_dir),inputs=net.pred_in,outputs={"prob_nms": net.pred_out["prob_nms"],"descriptors": net.pred_out["descriptors"],})

此时Batch还是1

后面onnx流程中就被固定住了

所以在：

pd转onnx的时候：

Dynamic Input Reshape Incorrect · Issue #1640 · onnx/tensorflow-onnx · GitHub

Dynamic Input Reshape Incorrect · Issue #8591 · microsoft/onnxruntime · GitHub