Torch Geometric Data Handling

data
1

Dear community,

I have problems with pytorch_geometric. I am training a graph neural network, and I encounter with a lot of problems. Firstly my Data object consist from x, edge_attr, edge_index, target and shortest_path_attr which is the shortest path matrix retrieved from Floyd-Warshall algorithm and its dimensionality is NxN where N is the number of nodes. When I try to give it to dataloader I have problems with Batching and simply tensors do not stack because of their difference of dimensions (the main problem comes from shortest path matrix). But when I pad the matrices and dataloader problem solves, I encounter another problem: my model doesn’t accept the batched tensor of X attributes (node features) because the dimensions of model and data do not match if I give a batch containing more then 1 graph info. What can I do, I would appreciate any help you can provide!

2

A bit late, but I just worked through this problem, also using the shortest-path distance matrix. What I did was just store the flattened version of each NxN matrix and then write a function for recovering the correct flattened index.

def batched_index_select(flattened_matrix, edge_index, batch):
    device = flattened_matrix.device
    sizes = torch.bincount(batch)

    batch_u = batch[edge_index[0]]

    node_offset = torch.cat(
        [torch.tensor([0]).to(device), torch.cumsum(sizes[:-1], dim=0)]
    )
    local_idx_u = edge_index[0] - node_offset[batch[edge_index[0]]]
    local_idx_v = edge_index[1] - node_offset[batch[edge_index[1]]]

    graph_offset = torch.cat(
        [torch.tensor([0]).to(device), torch.cumsum(sizes[:-1] ** 2, dim=0)]
    )

    flat_indices = graph_offset[batch_u] + sizes[batch_u] * local_idx_u + local_idx_v

    return flattened_matrix[flat_indices]