when converting model to onnx, Is there a way to specify output nodes

[tino926]

When exporting the model to ONNX, how can I obtain an output model with some final layers removed? To be more specific, I want to get an ONNX model (export from yolov5) with three output nodes, as marked in the image, and remove all layers afterward.

(請問在轉換訓練好的模型時至onnx時，是否能指定輸出的節點，並將其後的結點刪除呢？我發現最後幾層的reshape 等運算會拖慢某些AI accelerator，我都習慣另外自行處理其後的運算)

[Redreadblack]

我也遇到过这个问题。我的做法是用torch从新搭建一个模型，把后面的层删除，并将对应层的权重导入。
mmyolo中的image_demo.py中的init_detector可以帮助你构模型