Hi @albanD ,
Thanks for the idea! I didn’t realize you could access the input and output tensors that way. So your suggestion is to loop over all of the .modules() and register a hook for each of them. That would give me access to at least some of the tensors buried in the computation graph. And from there it would be easy to connect the tensors to the nodes in the .grad_fn computation graph I built using a dict. Restricting support to nn.Modules is certainly an option in the near term. Might be worth looking into.
I’m fairly certain that I will inevitably need to modify the PyTorch source in order to do exactly what I want. My goal is to find the path in the computation graph through which the most gradient flows for each each parameter. The tricky part is that, in order to do this correctly, the gradient needs to be known for every part of the computation graph, all the way from the loss node to the parameter nodes. The other tricky part is that every tensor operation (particularly matrix multiplications) encapsulates the equivalent of multiple scalar operations in the computation graph. So looking purely at the gradient tensors flowing in and out of each tensor operation is not enough – you need to have knowledge specific to each tensor operation to figure it out. That’s why I’m interested in accessing the forward value of each tensor operation – so I can recompute the scalar terms of the gradient and pick the ones with the highest magnitude.
Is there someone familiar with the source code who would know how to add the .output_tensor property I mentioned?
Thanks!