5.2 onnx推理显示shape 信息

前言

导出onnx过程中，有时会不显示shape信息，导致看起来不方便

步骤

下载必备包

pip install onnx
pip install onnx_graphsurgeon 

方式一

import onnx
onnx = onnx.load(onnx_path)
new_onnx= onnx.shape_inference.infer_shapes(onnx)
onnx.save(new_onnx, output_path)

方式二

import onnx
import onnx_graphsurgeon as gs
onnx = onnx.load(onnx_path)
graph = gs.import_onnx(onnx)
graph.cleanup().toposort() #从图形中删除未使用的节点和张量，并对图形进行拓扑排序

# Shape Estimation
new_onnx= None
try:
    new_onnx= onnx.shape_inference.infer_shapes(gs.export_onnx(graph))
except:
    new_onnx= gs.export_onnx(graph)
 onnx.save(new_onnx, output_path)