How to customly mini-batch HeteroData objects?

[DanielPerezJensen]

I am following along with the following tutorial about advanced mini batching.

I need to be able to do this because I am working with an LSTM within a graph structure. Basically, I want to use 6 graphs to predict 6 values for 4 of the nodes within a graph. To this end, I need a way to create a graph which contains 6 unconnected graphs. This seems simple enough when only using the normal Data for homogeneous graphs. But I am working in a heterogeneous setting.

I have 4 nodes for which I want to predict the following 6 values, and I have 11 nodes for which I have information that might be relevant to predicting these values. These types of nodes are different and have different features, also the edges between them are different.

There doesn't seem to be a guide on how to tackle this for the heterogeneous setting. The way I did it in the homogeneous setting is by basically holding a list for edge indices, a list for node attributes and the to be predicted values within a custom Data class called SeqData.

from torch_geometric.data import Data

class SeqData(Data):
    def __init__(self, edge_indices=None, xs=None, y=None):
        super().__init__()
        
        self.edge_indices = edge_indices
        self.xs = xs
        self.y = y
    
    def __inc__(self, key, value, *args, **kwargs):
        if key == "edge_indices":
            return self.xs[0].size(0)
        else:
            return super().__inc__(key, value, *args, **kwargs)

n_graphs = 2
n_nodes = 4

edge_index = torch.tensor([
    [1, 2, 3],
    [0, 0, 0]
])

edge_indices = edge_index.repeat(2, 1, 1)

xs = torch.stack([torch.randint(1, 5, (n_nodes, 8)) for _ in range(n_graphs)])
y = torch.randint(6, 10, (n_nodes, n_graphs))

data = SeqData(edge_indices, xs, y)

data_list = [data, data]

loader = DataLoader(data_list, batch_size=2, follow_batch=["xs", "y"])
batch = next(iter(loader))

Output:

SeqData(edge_indices=[2, 2, 3], xs=[2, 4, 8], y=[4, 2])
SeqDataBatch(edge_indices=[4, 2, 3], xs=[4, 4, 8], xs_batch=[4], y=[8, 2], y_batch=[8])

Keep in mind I am using random inputs for clarity's sake, so the number of graphs I store here are only 2, instead of the 6 I specified earlier. Also the number of nodes is different.

---

But how would I now do this with a HeteroData? When I try to do something similar like below:

from torch_geometric.data import Data, HeteroData

class SeqData(HeteroData):
    def __init__(self):
        super().__init__()
    
    def __inc__(self, key, value, store, *args, **kwargs):
        # not sure what to put here
        print(key, value)
        return 0
n_graphs = 2
n_nodes = 4

measurement_xs = torch.stack([torch.randint(1, 5, (n_nodes, 8)) for _ in range(n_graphs)])
measurement_y = torch.randint(6, 10, (n_nodes, n_graphs))

measurement_edge_index = torch.tensor([
    [1, 2, 3],
    [0, 0, 0]
])

measurement_edge_indices = edge_index.repeat(2, 1, 1)

data = SeqData()

data["measurement"].x = measurement_xs
data["measurement"].y = measurement_y
data["measurement", "flows", "measurement"].edge_index = measurement_edge_indices

data_list = [data, data]

loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))

print(batch)

Output:

SeqDataBatch()

I just get an empty batch, with no indication of what might be wrong. The __inc__ function doesn't even seem to be called. What else do I need to change to make this work?

[rusty1s]

This should be very similar to do for HeteroData. I am sorry this is not better documented. You should be able to do this by accessing the store inside __inc__ (not tested):

    def __inc__(self, key, value, store, *args, **kwargs):
        if key == "edge_indices":
            src_type, _, dst_type = store._key
            return torch.tensor([[0], [self[src_type.xs[0].size(0)]], [self[dst_type].xs[0].size(0)]])
        else:
            return super().__inc__(key, value, store, *args, **kwargs)

[DanielPerezJensen]

Thank you for your quick reply, I will look into it and report back. Maybe I can create an issue to add documentation about this?

[rusty1s]

Sure, please feel free to do so or directly send a PR that improves documentation in this regard. Thank you!