How to use batch adjacency matrix with GraphUNet

❓ Questions & Help

I want to directly input a batch of adjacency matrix(batch, n_nodes, n_nodes) and nodes(batch, n_nodes, feats) into the GraphUNet. I convert the dense matrix to sparse edge indexes. However the number of edges in the whole batch is various, the edge init edge_weight = x.new_ones(edge_index.size(1)) went wrong. Did I miss something for batch processing? How can I efficiently use the batch dense matrix as input?

[LingxiaoShawn]

Hi @rusty1s , how about this?

def to_sparse_batch(x, adj, mask):
    # transform x (B x N x D), adj (B x N x N), mask (B x N), here N is N_max
    # to x, edge_index, edge_attr/weight, batch
    B, N_max, D = x.shape
    # mask out adj 
    adj = (adj * mask.unsqueeze(2)).transpose(1,2) 
    adj = (adj * mask.unsqueeze(2)).transpose(1,2) 

    # get number of graphs 
    num_nodes_graphs = mask.sum(dim=1)  # B

    # get offset with size B
    offset_graphs = torch.cumsum(num_nodes_graphs, dim=0) # B
    
    # get x
    x = x.reshape(-1, D)[mask.reshape(-1)] # total_nodes * D
    
    # get weight and index
    edge_weight = adj[adj.nonzero(as_tuple=True)]
    nnz_index = adj.nonzero().t()
    graph_idx, edge_index = nnz_index[0], nnz_index[1:]

    # init batch
    batch = torch.zeros_like(x[:,0], dtype=torch.int64).fill_(B-1)

    # add offset to edge_index, and setup batch
    start = 0
    for i, offset in enumerate(offset_graphs[:-1]):
        edge_index[:, graph_idx==i+1] += offset
        batch[start:offset] = i
        start = offset

    return x, edge_index, edge_weight, batch

[rusty1s]

Sure, feel free to add it 😃