szagoruyko / pytorchviz

A small package to create visualizations of PyTorch execution graphs
MIT License
3.24k stars 279 forks source link

Lisibility of saved tensor #69

Open Varal7 opened 3 years ago

Varal7 commented 3 years ago

The commit https://github.com/pytorch/pytorch/commit/1efa863837aca33a611c36f030de9f30b7087ca4 changes the way variables are saved during the forward (in certain cases).

This can create comical graphs when visualized with torchviz because the saved variable shares the same id as the input.

For instance:

a = torch.randn(5, requires_grad=True)
y = a * a
make_dot(y, show_saved=True)

Before:

After https://github.com/pytorch/pytorch/commit/1efa863837aca33a611c36f030de9f30b7087ca4:

In this second image the blue node and the two orange nodes are merged into one.

I'm opening this issue to discuss how we would to fix this.

One option is to revert back to the old behavior by giving unique names to each node (so something like dot.node(SAVED_PREFIX + attr + str(id(val)), get_var_name(val, attr), fillcolor='orange')). The drawback is that the user loses the information that those three nodes are indeed backed by the same tensor.

Another option would be to draw a dotted edge between the saved variable and its base, for instance:

cc @soulitzer @albanD

albanD commented 3 years ago

From offline discussion, I think the dotted line version is nice. Note that this actually exposes some "internal" autograd stuff as these Tensors being the exact same python object or not is only true for non-outputs. Maybe we could change the way we do this detection and use tensor.data_ptr() to detect if SavedVariables ever match another Tensor already here. And use dotted lines between them (to specify they are views of each other).