DataLoader with variable sequence-lengths as target for Data objects

[DanielFPerez]

Hi all!

First, excuse in advance if this is not the category to place this issue, I will modify it if you point me to the correct one.

I have a supervised learning problem in which, for each graph, a sequence must be generated, whereby each element in the sequence is a vector representing a node classification.
Different graphs have different sequence lengths at the output.

e.g., here are some data samples:

As you see, the node features (Data.x) have size N x D (N nodes, D features), and the target vectors (Data.y) have size N x * (* is the sequence length for that specific graph.
The DataSet is a list of such samples.

the question is, under mini-batch training, how could I ensemble all the graphs with their sequence lengths using the DataLoader?
I have tried the standard way as described in the documentation, but I get a
RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 7 but got size 10 for tensor number 1 in the list.
So, a torch.cat error when trying to concatenate the different sequence lengths in torch_geometric/data/collate.py in _collate

Please let me know if there is a way of doing that please, otherwise I can proceed to modify the _collate function in the file and submit a pull request.

Thank you very much!

[rusty1s]

Yeah, the standard approach to solve this (as utilized in NLP) is to pad the label dimension before the torch.cat call. You can implement that as a transform on your data, e.g.:

def my_pad_transform(data):
    size = [data.y.size(0), max_seq - data.y.size(1)]
    data.y = torch.cat([data.y, data.y.new_zeros(size)], dim=-1)

Does that work for you?

[DanielFPerez]

Thanks for the quick response!

So you mean to call such a function right before the torch.cat in torch_geometric/data/collate.py::_collate rigth?

Yes, originally my intention was to pad the sequences, first I had planned to do it by overwriting the colalte_fn from the Dataloader, but then I realized torch_geometric::DataLoader already overwrites the collate_fn from PyTorch.

[rusty1s]

It depends on where you want to call it. You can call it on the full dataset prior to any DataLoader definition. This way, no modifications to DataLoader are needed but you might waste some memory. If you want to pad only across a mini-batch, you can write your own DataLoader:

def collate_fn(data_list):
    for data in data_list:
        # pad `data.y`
    return Batch.from_data_list(data)


class MyDataLoader(torch.utils.data.DataLoader):
    def __init__(self, *args, **kwargs):
        super().__init__(self, *args, collate_fn=collate_fn, **kwargs)

[DanielFPerez]

That looks perfect, it should work now :)
And thanks in general for your work on PyG, hope to keep working with it in the future!