Batched Graph Classification

Classify whole graphs. To train fast, put many small graphs into one batch, as graph libraries do.

Each graph is a pair (x, edge_index), with node features x of shape (n_i, f) and edges of shape (2, e_i). Row 0 of edge_index holds the source nodes and row 1 the target nodes. The task is to tell trees from graphs that contain a cycle. All node features are 1, so the model must use the structure.

Write three functions:

1. collate(graphs) joins a list of graphs into one big graph with no edges between the parts. Concatenate the features. Add an offset to each graph's node ids, equal to the number of nodes before it. Also return a batch vector that gives the graph index of each node.
2. sum_aggregate(x, edge_index) sends each source node's features to its target and adds them up: out[i] = sum of x[j] over edges j -> i. Use index_add. Do not build an adjacency matrix.
3. train_graph_classifier(graphs, labels) trains a GNN and returns a function that takes a list of graphs and returns their labels.

Function signature

def collate(graphs):                    # -> (x, edge_index, batch)
def sum_aggregate(x, edge_index):       # -> (n, f)
def train_graph_classifier(graphs, labels):   # -> predict(graphs)

Example

Input

g0 = (torch.ones(2, 1), torch.tensor([[0, 1], [1, 0]]))
g1 = (torch.ones(3, 1), torch.tensor([[0, 2], [2, 0]]))
collate([g0, g1])

Output

x          = ones of shape (5, 1)
edge_index = tensor([[0, 1, 2, 4],
                     [1, 0, 4, 2]])
batch      = tensor([0, 0, 1, 1, 1])

Reasoning

Graph 0 has 2 nodes, so graph 1's node ids move up by 2. Its edge 0 -> 2 becomes 2 -> 4. The first two nodes belong to graph 0 and the next three to graph 1.

Constraints and notes

• collate returns float features, an int64 edge_index of shape (2, E) and an int64 batch of shape (N,).
• sum_aggregate gives zeros for a node with no incoming edges. If an edge is listed twice, it counts twice. Gradients must flow to x.
• labels is an int64 tensor of 0 (tree) and 1 (has a cycle). The returned function must give int64 labels of shape (G,).
• Use sum pooling over the nodes of each graph, with index_add and the batch vector. Mean pooling loses the number of nodes.
• On held-out graphs, at least 90% of the predictions must be correct. A correct model gives the same predictions when the nodes of a graph are numbered in a different order.