How to extract layer shape and type from ONNX / PyTorch?

Question

I would like to 'translate' a PyTorch model to another framework (non-tf/keras).

I'm trying to take a pytorch model, and automate the translation to the other framework, which contains similar types of layers (i.e. conv2d, dense,...).

Is there a way from pytorch directly, or through onnx to retrieve a models layers, their types, shapes and connections ? (Weights are not important so far)

Answer 1

From discussion in comments on your question:

each node in onnx has a list of named inputs and a list of named outputs.

For the input list accessed with node.input you have for each input index either the graph input_name that feeds that input or the name of a previous output that feeds that input. There are also initializers which are onnx parameters.

# model is an onnx model
graph = model.graph
# graph inputs
for input_name in graph.input:
    print(input_name)
# graph parameters
for init in graph.init:
    print(init.name)
# graph outputs
for output_name in graph.output:
    print(output_name)
# iterate over nodes
for node in graph.node:
    # node inputs
    for idx, node_input_name in enumerate(node.input):
        print(idx, node_input_name)
    # node outputs
    for idx, node_output_name in enumerate(node.output):
        print(idx, node_output_name)

Shape inference is talked about here and for python here

The gist for python is found here

Reproducing the gist from 3:

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

and find the shape info in inferred_model.graph.value_info.

You can also use netron or from GitHub to have a visual representation of that information.