FM / data

/usr/local/lib64/python3.6/site-packages/torch/utils/data/

/usr/local/lib64/python3.6/site-packages/torch/utils/data

Name	Size	Mode	Actions
communication/	-	0755	rm
datapipes/	-	0755	rm
_utils/	-	0755	rm
__pycache__/	-	0755	rm
backward_compatibility.py	254	0644	edit dl rm
dataloader.py	67104	0644	edit dl rm
dataloader_experimental.py	6995	0644	edit dl rm
dataset.py	15001	0644	edit dl rm
distributed.py	5689	0644	edit dl rm
graph.py	1115	0644	edit dl rm
sampler.py	10226	0644	edit dl rm
sharding.py	1079	0644	edit dl rm
_decorator.py	7699	0644	edit dl rm
_typing.py	14436	0644	edit dl rm
__init__.py	2070	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/utils/data/graph.py (1115B)

import io
import pickle

from torch.utils.data import IterableDataset

from typing import Any, Dict

reduce_ex_hook = None

def stub_unpickler():
    return "STUB"

# TODO(VitalyFedyunin): Make sure it works without dill module installed
def list_connected_datapipes(scan_obj):

f = io.BytesIO()
    p = pickle.Pickler(f)  # Not going to work for lambdas, but dill infinite loops on typing and can't be used as is

def stub_pickler(obj):
        return stub_unpickler, ()

captured_connections = []

def reduce_hook(obj):
        if obj == scan_obj:
            raise NotImplementedError
        else:
            captured_connections.append(obj)
            return stub_unpickler, ()

# TODO(VitalyFedyunin):  Better do it as `with` context for safety
    IterableDataset.set_reduce_ex_hook(reduce_hook)
    p.dump(scan_obj)
    IterableDataset.set_reduce_ex_hook(None)
    return captured_connections

def traverse(datapipe):
    items = list_connected_datapipes(datapipe)
    d: Dict[Any, Any] = {datapipe: {}}
    for item in items:
        d[datapipe].update(traverse(item))
    return d

import io import pickle from torch.utils.data import IterableDataset from typing import Any, Dict reduce_ex_hook = None def stub_unpickler(): return "STUB" # TODO(VitalyFedyunin): Make sure it works without dill module installed def list_connected_datapipes(scan_obj): f = io.BytesIO() p = pickle.Pickler(f) # Not going to work for lambdas, but dill infinite loops on typing and can't be used as is def stub_pickler(obj): return stub_unpickler, () captured_connections = [] def reduce_hook(obj): if obj == scan_obj: raise NotImplementedError else: captured_connections.append(obj) return stub_unpickler, () # TODO(VitalyFedyunin): Better do it as `with` context for safety IterableDataset.set_reduce_ex_hook(reduce_hook) p.dump(scan_obj) IterableDataset.set_reduce_ex_hook(None) return captured_connections def traverse(datapipe): items = list_connected_datapipes(datapipe) d: Dict[Any, Any] = {datapipe: {}} for item in items: d[datapipe].update(traverse(item)) return d