While that works, it also loses kernel size information from parameters to be included in the graph. Alternatively, you can follow the error and change __getitem__ to _node_get. To do that, include the following in the beginning of your pytorch_builder.py
# From https://github.com/pytorch/pytorch/blob/2efe4d809fdc94501fc38bf429e9a8d4205b51b6/torch/utils/tensorboard/_pytorch_graph.py#L384
def _node_get(node: torch._C.Node, key: str):
    """Gets attributes of a node which is polymorphic over return type."""
    sel = node.kindOf(key)
    return getattr(node, sel)(key)

torch._C.Node.__getitem__ = _node_get
Now, you should be able to extract from the nodes just fine, so there should be no issues with this line.
params = {k: torch_node[k] for k in torch_node.attributeNames()}
Credit: TypeError: 'torch._C.Node' object is not subscriptable
still do not show data flow dimension

modify pytorch_builder.py get_shape(torch_node) method as follows:

def get_shape(torch_node):
    """Return the output shape of the given Pytorch node."""
    # Extract node output shape from the node string representation
    # This is a hack because there doesn't seem to be an official way to do it.
    # See my quesiton in the PyTorch forum:
    # https://discuss.pytorch.org/t/node-output-shape-from-trace-graph/24351/2
    # TODO: find a better way to extract output shape
    # TODO: Assuming the node has one output. Update if we encounter a multi-output node.

    # m = re.match(r".*Float\(([\d\s\,]+)\).*", str(next(torch_node.outputs())))
    # if m:
    #     shape = m.group(1)
    #     shape = shape.split(",")
    #     shape = tuple(map(int, shape))
    # else:
    #     shape = None
    shape = next(torch_node.outputs()).type().sizes()
    return shape

AttributeError: module 'torch.onnx' has no attribute '_optimize_trace' #100

Description

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions