Data Prefetching/Pre-loading for Graph Data

[FarzanT]

Hello,
Tensor prefetching using Pytorch is relatively simple, and is done by something like:

class DataPrefetcher():
    def __init__(self, loader):
        self.loader = iter(loader)
        self.stream = torch.cuda.Stream()
        self.preload()

    def __iter__(self):
        return self

    def preload(self):
        try:
            self.next_input, self.next_target, self.next_weight = next(self.loader)
        except StopIteration:
            self.next_input = None
            self.next_target = None
            return
        with torch.cuda.stream(self.stream):
            # Convert to float then transfer to GPU non-blocking-ly
            self.next_input = [enc_data.float() for enc_data in self.next_input]
            self.next_target = self.next_target.float()

            self.next_input = [enc_data.cuda(non_blocking=True) for enc_data in self.next_input]
            self.next_target = self.next_target.cuda(non_blocking=True)

    def __next__(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        input = self.next_input
        target = self.next_target
        weight = self.next_weight
        if input is not None:
            for i in input:
                i.record_stream(torch.cuda.current_stream())
        if target is not None:
            target.record_stream(torch.cuda.current_stream())
        if weight is not None:
            weight.record_stream(torch.cuda.current_stream())

        self.preload()
        return input, target, weight

Here, upcoming batches of data are put on a cuda stream and are pre-loaded on the GPU while the current batch is returned.
This makes use of torch.Tensor.record_stream() (https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) but PyG graph data does not have such a method. Does anyone know an alternative for pre-loading graph data onto the GPU?

Thanks!

[rusty1s]

Thanks for starting the discussion. This is something which I wanted to add to PyG long time ago. For allowing record_stream functionality, one can iterate over data and call record_stream on every Tensor attribute:

for key, value in data:
    if isinstance(value, Tensor):
        value.record_stream(...)

Let me know if you have any interest in contributing this feature.

[FarzanT]

Thank you Matthias, I'll make a pull request shortly!