Add graphs of custom computation graph to tensorboard

I was wondering if there is a way to add a custom computation graph to tensorboard using PyTorch. The website says it works for PyTorch models and tensors. I didn’t find examples where the graphs are being added to tensorboard for custom computation graphs.

So basically my function takes in a couple of tensors, performs calculations within the function creating more tensors and then outputs a result. I wanted to be able to visualize the interactions between the various tensors.

I tried calling writer.add_graph(my_func, (my_args1, my_args2)), but that resulted in an error stating : 'function' object has no attribute 'param_init_net'

I have tried to create an extremely simple working example here.
https://colab.research.google.com/drive/1AIcObBl1QTc3h0hryWo1ZO7ZJ5hk_MfT?usp=sharing

Can someone help me out on this?

Thank you.

1 Like

@pecey Did you have any luck with this issue?

I am also trying to get a graph visualisation for arbitrary Pytorch code/functions. To me, it seems that add_graph() requires a model as an input. I am using torch for autograd functionality without DL. Is it possible to get a computational graph like in TF?