add utitlity to visualize datapipe graphs

[pmeier]

Closes #299.

[ejguan]

Noob question: Wondering do we have a potential extensibility using graphviz to visualize the attributes for each Node?

[pmeier]

Wondering do we have a potential extensibility using graphviz to visualize the attributes for each Node?

Sure. Everything we put in Node.__str__ will be displayed. For example, we could add the name of the functions of a Filter or Mapper to the node. Doing this would turn

from torchdata.datapipes.iter import IterableWrapper


def no_map(x):
    return x


def no_filter(x):
    return True


dp1, dp2, dp3 = IterableWrapper([]).fork(3)
dp1 = dp1.map(no_map)
dp2 = dp2.filter(no_filter)
dp = dp1.zip(dp2, dp3)

into

This is why I suggested #284 (comment). This would reduce Node.__str__ to return str(self.dp).

[pmeier]

@ejguan in #237 (comment)

Wondering what would be the visualized graph with such circular reference?

from torchdata.datapipes.iter import IterDataPipe, IterableWrapper

class CustomIterDataPipe(IterDataPipe):
    def noop(self, x):
        return x + 1

    def __init__(self):
        self._dp = IterableWrapper([1, 2, 4]).map(self.noop)

    def __iter__(self):
        yield from self._dp

Due to the circular reference, we either need to use traverse(..., only_datapipe=True) as suggested in #330 (comment) or fix the traversal similar to what I proposed in #237 (comment). Otherwise, I don't think we need extra handling for "circular" datapipes.

[ejguan]

This looks good to me.
@NivekT What do you think? We have inline doc for to_graph. Do we need an example for users shown in README or another place?

[NivekT]

Sorry for not getting back to this earlier. Overall, it looks very good to me!

I tried playing around with it to see if there are cases where the output is not sensible. This is the only case where things may be confusing, and I think we may want to show the ChildDataPipe (perhaps renamed to something else like DemuxChild1 or something). Admittedly, I think the code snippet below may not be representative of the usual use cases of DataPipes in a pipeline:

from torchdata.datapipes.iter import IterableWrapper
from torchdata.datapipes.utils import to_graph

dp = IterableWrapper(range(10))
dp = dp.map(lambda x: x + 1)
dp = dp.filter(lambda _: True)
dp1, dp2 = dp.demux(num_instances=2, classifier_fn=lambda x: x % 0 == 0)
dp3 = dp1.zip(dp2)
g = to_graph(dp3)
g.view()

We can also apply the function on a few torchvision data pipelines and see what they look like.

[pmeier]

I wrote a small script to test this on the torchvision datasets. Assuming this PR is merged and you have a nightly from torchdata and torchvision installed, this is the result.

very easy

CIFAR10

easy

MNIST

medium

VOC

difficult

COCO

CelebA

[NivekT]

I wrote a small script to test this on the torchvision datasets. Assuming this PR is merged and you have a nightly from torchdata and torchvision installed, this is the result.

very easy

CIFAR10

easy

MNIST

medium

VOC

difficult

COCO
CelebA

I wrote a small script to test this on the torchvision datasets. Assuming this PR is merged and you have a nightly from torchdata and torchvision installed, this is the result.

very easy

CIFAR10

easy

MNIST

medium

VOC

difficult

COCO
CelebA

Those graphs look complicated but sensible to me. Unless there are ways to simplify the pre-processing logic, I think they look good. I think the remaining part is the edge case with _ChildDataPipes combining back together immediately.

[pmeier]

@NivekT I've added a debug: bool = False flag as explained in #330 (comment). This gives the following results for your example in #330 (review):

debug=False (default)

debug=True

`

[NivekT]

This PR is good to land as it is. I will follow up with a different PR of documentations.

[facebook-github-bot]

@NivekT has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.