comaniac commented on a change in pull request #7055:
URL: https://github.com/apache/tvm/pull/7055#discussion_r538658159
##########
File path: python/tvm/driver/tvmc/frontends.py
##########
@@ -161,14 +161,11 @@ def load(self, path):
# pylint: disable=E1101
model = onnx.load(path)
+ shape_dict = {}
# pylint: disable=E1101
- name = model.graph.input[0].name
-
- # pylint: disable=E1101
- proto_shape = model.graph.input[0].type.tensor_type.shape.dim
- shape = [d.dim_value for d in proto_shape]
-
- shape_dict = {name: shape}
+ for i in model.graph.input:
+ # pylint: disable=E1101
+ shape_dict[i.name] = [dim.dim_value for dim in
i.type.tensor_type.shape.dim]
Review comment:
I agree with @leandron. It would be better to keep TVMC as a command
line argument processor and put other logic to other modules as we can. The
logic @dlexplorer proposed looks good to me. In addition, when the shape param
is passed, we could also check if it is consistent with the one defined in the
model, or throw out an error or a warning otherwise. In this case, we don't
even need a test case in this PR.
Also cc @mbrookhart who recently works on the ONNX frontend.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
