onnx图，如何获取输出维度？

Question

如何获取onnx神经网络输出层的维度？

我可以得到 onnx 图，但没有输出尺寸：

~/onnx-tensorrt/third_party/onnx/onnx/tools/net_drawer.py --input ./weights/tiny_3l_v5_11_608.onnx  --output ./weights/tiny_3l_v5_11_608.dot --embed_docstring

特尔维辛，马库斯

Answer 1

你可以简单地使用

from onnx import shape_inference
inferred_model = shape_inference.infer_shapes(original_model)

并在 中找到形状信息inferred_model.graph.value_info。

有关完整文档，请参阅https://github.com/onnx/onnx/blob/master/docs/ShapeInference.md

Answer 2

嘿，

https://github.com/onnx/onnx-tensorrt 有这个功能。

我做到了

onnx2trt yolo_tiny_3l.onnx  -t yolo_tiny_3l.onnx.txt -l -v

特尔维辛，马库斯