Batching a list of lists with different length

[zhaoyu-li]

Hi,

I encounter a problem when using torch_geometric. Below is an example: indices is a list of lists with different lengths. What I actually want to get after batching is [[[0, 1, 2, 3], [2, 4], [5, 6]], [[7, 8, 9, 10], [9, 11]]], however, after batching, I get: [[[0, 1, 2, 3], [2, 4], [5, 6]], [[0, 1, 2, 3], [2, 4]]].

I also try to make the elements in the list be the long tenor but it doesn't work because of the different lengths. I also see the implementation in data/collate.py, it seems like there is no support for an operation like this. Could you please help me find a way to solve this?

Thanks in advance!

import torch
import torch.nn as nn
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader


class MyData(Data):
    def __init__(self, x, indices):
        super().__init__()
        self.x = x
        self.indices = indices
    
    def __inc__(self, key, value, *args, **kwargs):
        if key == 'indices':
            return self.x.size(0)
        else:
            return super().__inc__(key, value, *args, **kwargs)


x = torch.randn(7, 1)
indices = [[0, 1, 2, 3], [2, 4], [5, 6]]

d = MyData(x, indices)
data_list = [d, d]
loader = DataLoader(data_list, batch_size=2)
print(next(iter(loader)).indices)

[rusty1s]

Yeah, you are right. Data.__inc__ is currently only applied to PyTorch tensors. We probably need to at least omit a warning in case a user wants to increment a non-incremental object. Let me know if you have interest in adding this :)

One workaround is to save indices as a one-dimensional Tensor and to utilize torch.Tensor.split after batching:

import torch
import torch.nn as nn
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader


class MyData(Data):
    def __init__(self, x, indices):
        super().__init__()
        self.x = x
        self.indices = indices
    
    def __inc__(self, key, value, *args, **kwargs):
        if key == 'indices':
            return self.x.size(0)
        else:
            return super().__inc__(key, value, *args, **kwargs)


x = torch.randn(7, 1)
indices = [[0, 1, 2, 3], [2, 4], [5, 6]]

d = MyData(x, indices)
data_list = [d, d]
loader = DataLoader(data_list, batch_size=2)
print(next(iter(loader)).indices)

This should also be way faster than working with Python lists.

[zhaoyu-li]

Thanks a lot!