How to create labels with adjacency matrix shape

[vietngth]

Hi, just a follow-up for #8548 2) point.

As my final task is to predict a label matrix with shape [nodes, nodes] for each graph, what would be a way to store these labels correctly? For example, 2 Data objects has the following format:

Data(x=[231, 10], edge_index=[2, 988], edge_attr=[988, 7], y=[1, 231, 231])
Data(x=[41, 10], edge_index=[2, 92], edge_attr=[92, 7], y=[1, 41, 41])

When using batch for training, the following error occured:

torch_geometric\data\collate.py:177: RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 231 but got size 41 for tensor number 1 in the list.

As suggested in the previous post, they should be stored as [1, max_nodes, max_nodes], where max_nodes is the largest number of nodes of the graph in the current batch. However, this error occurs during the loading batch process before the forward pass. Is there a way to access max_nodes of each batch to store the labels in the correct shapes?

I have a custom class of dataset: class CustomDataset(torch_geometric.data.Dataset) for some custom logic of data creation, and CustomDataModule to handle data loaders logic for torch.geometric.loader.DataLoader:

class CadDatamodule(lightning.LightningDataModule):
    def __init__(self, **kargs):

    def setup(self, stage: str):

    def train_dataloader(self):
        return DataLoader(...)
    ...

[vietngth]

Update: it seems the user input collate_fn is always discarded in torch_geometric.loader.DataLoader as stated here: #3426.

I have tried to overcome this by the following:

1. Copy and paste the whole DataLoader from here: https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/loader/dataloader.py#L71
2. Directly adjust the padding logic inside collate_fn:

class Collater:
    ...
    def collate_fn(self, batch: List[Any]) -> Any:
        if isinstance(self.dataset, OnDiskDataset):
            return self(self.dataset.multi_get(batch))

        # additional padding logic
        max_num_nodes = max(graph.x.shape[0] for graph in batch)
        new_batch = []
        for graph in batch:
            num_nodes = graph.x.shape[0]
            pad_slots = max_num_nodes - num_nodes
            padded_label = pad(
                input=graph.label,
                pad=(0, pad_slots, 0, pad_slots),
                mode="constant",
                value=0
            )
            padded_graph = Data(
                x=graph.x,
                edge_index=graph.edge_index,
                edge_attr=graph.edge_attr,
                part=graph.part,
                label=padded_label
            )
            new_batch.append(padded_graph)

        return self(new_batch)

This work-around works as intended, but not really elegant with the copy-paste. Any better ideas would be appreciate!

[rusty1s]

Yes, this is the intended workaround for this. Do you have an idea to make this more elegant?

[vietngth]

I don't think I can think of anything at the moment, but will keep this in mind in the future. I'll close this one, thank you for clarification!