(
self,
graph,
indices,
graph_sampler,
device=None,
use_ddp=False,
ddp_seed=0,
batch_size=1,
drop_last=False,
shuffle=False,
use_prefetch_thread=None,
use_alternate_streams=None,
pin_prefetcher=None,
use_uva=False,
gpu_cache=None,
**kwargs,
)
| 923 | """ |
| 924 | |
| 925 | def __init__( |
| 926 | self, |
| 927 | graph, |
| 928 | indices, |
| 929 | graph_sampler, |
| 930 | device=None, |
| 931 | use_ddp=False, |
| 932 | ddp_seed=0, |
| 933 | batch_size=1, |
| 934 | drop_last=False, |
| 935 | shuffle=False, |
| 936 | use_prefetch_thread=None, |
| 937 | use_alternate_streams=None, |
| 938 | pin_prefetcher=None, |
| 939 | use_uva=False, |
| 940 | gpu_cache=None, |
| 941 | **kwargs, |
| 942 | ): |
| 943 | # (BarclayII) PyTorch Lightning sometimes will recreate a DataLoader from an existing |
| 944 | # DataLoader with modifications to the original arguments. The arguments are retrieved |
| 945 | # from the attributes with the same name, and because we change certain arguments |
| 946 | # when calling super().__init__() (e.g. batch_size attribute is None even if the |
| 947 | # batch_size argument is not, so the next DataLoader's batch_size argument will be |
| 948 | # None), we cannot reinitialize the DataLoader with attributes from the previous |
| 949 | # DataLoader directly. |
| 950 | # A workaround is to check whether "collate_fn" appears in kwargs. If "collate_fn" |
| 951 | # is indeed in kwargs and it's already a CollateWrapper object, we can assume that |
| 952 | # the arguments come from a previously created DGL DataLoader, and directly initialize |
| 953 | # the new DataLoader from kwargs without any changes. |
| 954 | if isinstance(kwargs.get("collate_fn", None), CollateWrapper): |
| 955 | assert batch_size is None # must be None |
| 956 | # restore attributes |
| 957 | self.graph = graph |
| 958 | self.indices = indices |
| 959 | self.graph_sampler = graph_sampler |
| 960 | self.device = device |
| 961 | self.use_ddp = use_ddp |
| 962 | self.ddp_seed = ddp_seed |
| 963 | self.shuffle = shuffle |
| 964 | self.drop_last = drop_last |
| 965 | self.use_prefetch_thread = use_prefetch_thread |
| 966 | self.use_alternate_streams = use_alternate_streams |
| 967 | self.pin_prefetcher = pin_prefetcher |
| 968 | self.use_uva = use_uva |
| 969 | kwargs["batch_size"] = None |
| 970 | super().__init__(**kwargs) |
| 971 | return |
| 972 | |
| 973 | # (BarclayII) I hoped that pin_prefetcher can be merged into PyTorch's native |
| 974 | # pin_memory argument. But our neighbor samplers and subgraph samplers |
| 975 | # return indices, which could be CUDA tensors (e.g. during UVA sampling) |
| 976 | # hence cannot be pinned. PyTorch's native pin memory thread does not ignore |
| 977 | # CUDA tensors when pinning and will crash. To enable pin memory for prefetching |
| 978 | # features and disable pin memory for sampler's return value, I had to use |
| 979 | # a different argument. Of course I could change the meaning of pin_memory |
| 980 | # to pinning prefetched features and disable pin memory for sampler's returns |
| 981 | # no matter what, but I doubt if it's reasonable. |
| 982 | self.graph = graph |
no test coverage detected