demonstrate how to trace/parse a pytorch graph

parse_pytorch_graph.py

import torch
from torchvision import models

def parse(net, inputs = torch.randn(1, 3, 224, 224)):
    with torch.onnx.set_training(net, False):
        trace = torch.onnx.utils._trace(net, inputs)

    graph = trace.graph()
    for n in graph.nodes():
        print(n.scopeName(), n.kind())
        attrs = str({k: n[k] for k in n.attributeNames()})
        print("attrs", attrs)
        inputs = [i.uniqueName() for i in n.inputs()]
        print("inputs", inputs)
        outputs = [i.uniqueName() for i in n.outputs()]
        print("outputs", outputs)
        print('---------------\n\n')

parse(models.densenet121(pretrained=False))