Obtaining unit-wise forward computational graph

[jordiae]

Hi, thanks for your work.

One question: do you know whether it would be feasible to extract the unit-level, forward computational graph of a PyTorch model using some of the internal APIs of tensorboardX? Not necessarily to visualize it (which would be messy), but, for instance, to encode the computational graph of the network into an adjacency matrix for research purposes.

Many thanks in advance.

[lanpa]

Hi there is not such API. But you might want to have a look at: https://github.com/lanpa/tensorboardX/blob/master/tensorboardX/pytorch_graph.py#L293
The link of a node will be iterated there. I think it is possible to identify the and map those nodes by the id() function.
Hope that helps. Note that pytorch_graph.py is legacy (the current tensorboardX relays the graph creation job directly to pytorch), so the current version of pytorch may not work directly.