Extracting execution DAG from PyTorch Module
Let us consider a simple feed-forward deep net or even a resnet. In both case, you can think of the execution flow as a DAG (directed acyclic graph) where nodes are layers (it's a chain if there's no skip connexion).
For some reason, I need this DAG for pyTorch models. Up to now I have been using custom object subclassing nn.Module (that stores the information that I need) but I would like to be able ultimately to use an arbitrary model in input of my pipeline. Is there a way to extract this DAG automatically ? (it does not matter if it's not working in all cases, I am mainly looking for general ideas / principles to achieve my goal).

That is indeed possible. Consider following evaluation of a simple formula as a stand in for a more complicated network:
import torch
a = torch.randn(3, requires_grad=True)
b = 3*a + 1
c = torch.relu(b)
d = c.sum()

Then d has a .grad_fn attribute, from which you can recurse through the evaluation graph like so
d.grad_fn
d.grad_fn.next_functions
d.grad_fn.next_functions
d.grad_fn.next_functions[0][0].next_functions
d.grad_fn.next_functions[0][0].next_functions[0][0].next_functions

Basically next_functions gives you a list of the arguments to the current module/operation, and each entries of this list is a tuple with the actual object, as well as an integer that indicates the position of the argument. This is documented in more detail here.
If you don't want to do it yourself, you can also use the visualization tool for the computation graph that is built into tensorboard, that is SummaryWriter.add_graph() as documented here.
