Loading TemporalData into DataLoader

[yulonglin]

I have been using TemporalData objects to load JODIE-formatted temporal data: https://github.com/rusty1s/pytorch_geometric/blob/4d062ae84e1c0353e7c8d15481f8b6e6d3f2ecca/torch_geometric/data/temporal.py

However, TemporalData is a slightly odd datatype, so I've written my own custom code so I can wrap it in PyTorch DataLoader objects for parallelisation (appended below).

Does it make sense for me to do a PR? Most of what the following code does can be covered by using a TemporalDataset class to wrap around TemporalData, and adding a case in the Collater class, but I'm not sure it belongs alongside the other cases: https://github.com/rusty1s/pytorch_geometric/blob/4d062ae84e1c0353e7c8d15481f8b6e6d3f2ecca/torch_geometric/data/dataloader.py#L13

import torch
from torch.utils.data.dataloader import default_collate
from torch_geometric.data import TemporalData


class TemporalDataset(torch.utils.data.Dataset):
    def __init__(self, temporal_data: TemporalData):
        self.temporal_data = temporal_data

    def __len__(self):
        return len(self.temporal_data.src)

    def __getitem__(self, idx):
        return self.temporal_data[idx]


class TemporalDataCollater(object):
    def collate(self, temporal_data_list):
        batch = TemporalData()
        for key in temporal_data_list[0].keys:
            batch[key] = default_collate([d[key] for d in temporal_data_list
                                          ]).squeeze(1)
        return batch

    def __call__(self, batch):
        return self.collate(batch)


class TemporalDataLoader(torch.utils.data.DataLoader):
    def __init__(self,
                 dataset: TemporalDataset,
                 batch_size=1,
                 shuffle=False,
                 **kwargs):

        if "collate_fn" in kwargs:
            del kwargs["collate_fn"]

        super(TemporalDataLoader,
              self).__init__(dataset,
                             batch_size,
                             shuffle,
                             collate_fn=TemporalDataCollater(),
                             **kwargs)

[rusty1s]

This is great. Please feel free to add it :)

[yulonglin]

Added it as a PR (#2625) :)