Generate graph visualization for torchvision datasets

main.py

import pathlib
import warnings

import tqdm

from torchdata.datapipes.utils import to_graph
from torchvision.prototype import datasets

directory = pathlib.Path.cwd() / "torchvision_datasets"
directory.mkdir(parents=True, exist_ok=True)

format = "png"

manual_configs = {
    "clevr": [
        dict(split="train"),
        dict(split="test"),
    ],
    "coco": [
        dict(annotations="instances"),
        dict(annotations=None),
    ],
    "cub200": [
        dict(year="2010"),
        dict(year="2011"),
    ],
    "gtsrb": [
        dict(split="train"),
        dict(split="test"),
    ],
    "imagenet": [
        dict(split="train"),
        dict(split="val"),
        dict(split="test"),
    ],
    "sbd": [
        dict(split="train"),
        dict(split="train_noval"),
    ],
    "stanford-cars": [
        dict(split="train"),
        dict(split="test"),
    ],
}

for name, configs in tqdm.tqdm([(name, manual_configs.get(name, [dict()])) for name in datasets.list_datasets()]):
    for config in tqdm.tqdm(configs):
        stem = "_".join([name, *[f"{k}_{v}" for k, v in sorted(config.items())]])
        if (directory / f"{stem}.{format}").exists():
            continue

        try:
            dataset = datasets.load(name, **config)
        except Exception as error:
            warnings.warn(f"Dataset '{name}' with config {config} raised the following exception:\n\n{error}")
            continue

        graph = to_graph(dataset)

        graph.render(filename=stem, directory=directory, format=format, cleanup=True)