search

ThanatosShinji / onnx-tool

A parser, editor and profiler tool for ONNX models.
https://pypi.org/project/onnx-tool/
MIT License
352 stars 44 forks source link

model_profile fails to run on onnx model with 4 dimensional input #31

Closed hamoudyounesyounes closed 1 year ago

hamoudyounesyounes commented 1 year ago

I have a pytorch model of UNet 3D (https://github.com/wolny/pytorch-3dunet), it is successfully exported to onnx by supplyinh an input _3d_x = torch.rand(1, 1, 32, 64, 64) and the graph looks fine. However, when I call:

nnx_tool.model_profile('./onnx_models/UNet3D.onnx', savenode='node_table.txt') an error is issued:

line 84, in model_profile g.shape_infer(dynamic_shapes) lib/python3.10/site-packages/onnx_tool/graph.py", line 828, in shape_infer oshapes = node.shape_infer(itensors) lib/python3.10/site-packages/onnx_tool/node.py", line 748, in shape_infer outshape[self.axis] += shape[self.axis] IndexError: list index out of range

does onnx_tool.model_profile( ) supports inputs as (Batch, depth, ch, h, w) ?

ThanatosShinji commented 1 year ago

Yes, it supports. And onnx_tool tested one 3D Segmentation model before. Can you provide some script to export the UNet3D.onnx from your repo?

ThanatosShinji commented 1 year ago

Never mind. I use --config resources/3DUnet_lightsheet_nuclei/test_config.yaml to reproduce your issue. The bug code is fixed in the main branch. The pypi package will be released after model UTs. You can install the package through the repo to validate the fix.

hamoudyounesyounes commented 1 year ago

I can confirm it is passing now. Thanks.