A minimal graph

Now, a basic Data object is created from three tensors. The first holds the node features: two features each for five nodes. (Both features are identical on purpose, for “cognitive ease” – on our, not the algorithm’s, part.)

import torch

x = torch.tensor([[1, 1], [2, 2], [3, 3], [11, 11], [12, 12]], dtype=torch.float)

The second specifies existing connections. For undirected graphs (like ours), each edge appears twice. The tensor you see here is specified in one-edge-per-line form for convenience reasons; to the Data() constructor we’ll pass its transpose instead.

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

The third tensor holds the node labels. (The task will be one of node – not edge, not graph – classification.)

y = torch.tensor([[0], [0], [0], [1], [1]], dtype=torch.float)

Constructing and inspecting the resulting graph, we have:

from torch_geometric.data import Data

data = Data(x = x, edge_index = edge_index.t().contiguous(), y = y)
data.x
data.edge_index
data.y

tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 3.,  3.],
        [11., 11.],
        [12., 12.]])
        
tensor([[0, 1, 0, 2, 1, 2, 2, 3, 2, 4, 3, 4],
        [1, 0, 2, 0, 2, 1, 3, 2, 4, 2, 4, 3]])
        
tensor([[0.],
        [0.],
        [0.],
        [1.],
        [1.]])

For our upcoming experiments, it’s more helpful, though, to visualize the graph:

import networkx as nx
import matplotlib.pyplot as plt
from torch_geometric.utils import to_networkx

def visualize_graph(G, color, labels):
    plt.figure(figsize=(7,7))
    plt.axis('off')
    nx.draw_networkx(
      G,
      pos = nx.spring_layout(G, seed = 777),
      labels = labels,
      node_color = color,
      cmap = "Set3"
      )
    plt.show()

G = to_networkx(data, to_undirected = True, node_attrs = ["x"])
labels = nx.get_node_attributes(G, "x")
visualize_graph(G, color = data.y, labels = labels)

Although our experiments won’t be about training performance (how could they be, with just five nodes), let me remark in passing that this graph is small, but not boring: The middle node is equally connected to both “sides”, yet feature-wise, it would pretty clearly appear to belong on just one of them. (Which is true, given the provided class labels). Such a constellation is interesting because, in the majority of networks, edges indicate similarity.

A minimal GNN

Now, we code and run the minimal GNN. We’re not interested in class labels (yet); we just want to see each node’s embeddings after a single pass.

from torch_geometric.nn import MessagePassing

class IAmLazy(MessagePassing):
    def forward(self, x, edge_index):
        out = self.propagate(edge_index, x = x)
        return out
      
module = IAmLazy()
out = module(data.x, data.edge_index)
out

tensor([[ 5.,  5.],
        [ 4.,  4.],
        [26., 26.],
        [15., 15.],
        [14., 14.]])

Evidently, we just had to start the process – but what process, exactly? From what we know about the three stages of message passing, an essential question is what nodes do with the information that flows over the edges. Our first experiment, then, is to inspect the incoming messages.

Poking into message()

In message(), we have access to a structure named x_j. This tensor holds, for each node , the embeddings of all nodes connected to it via incoming edges. We’ll print them, and then, just return them, unchanged.

class IAmMyOthers(MessagePassing):
    def forward(self, x, edge_index):
        out = self.propagate(edge_index, x = x)
        return out
    def message(self, x_j):
        print("in message, x_j is")
        print(x_j)
        return x_j
      
module = IAmMyOthers()
out = module(data.x, data.edge_index)
print("result is:")
out

in message, x_j is
tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 1.,  1.],
        [ 3.,  3.],
        [ 2.,  2.],
        [ 3.,  3.],
        [ 3.,  3.],
        [11., 11.],
        [ 3.,  3.],
        [12., 12.],
        [11., 11.],
        [12., 12.]])
        
result is:
tensor([[ 5.,  5.],
        [ 4.,  4.],
        [26., 26.],
        [15., 15.],
        [14., 14.]])

Let me spell this out. In data.edge_index, repeated here for convenience:

tensor([[0, 1, 0, 2, 1, 2, 2, 3, 2, 4, 3, 4],
        [1, 0, 2, 0, 2, 1, 3, 2, 4, 2, 4, 3]])

the first pair denotes the edge from node 0 (that had features (1, 1)) to node 1. This information is found in x_j’s first row. Then the second row holds the information flowing in the opposite direction, namely, the features associated with node 1. And so on.

Interestingly, since we’re passing through this module just once, we can see the messages that will be sent without even running it.

Namely, since data.edge_index[0] designates the source nodes for each edge:

data.edge_index[0]

we can index into data.x to pick up what will be the incoming features for each connection.

data.x[data.edge_index[0]]

tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 1.,  1.],
        [ 3.,  3.],
        [ 2.,  2.],
        [ 3.,  3.],
        [ 3.,  3.],
        [11., 11.],
        [ 3.,  3.],
        [12., 12.],
        [11., 11.],
        [12., 12.]])

Now, what does this tell us? Node 0, for example, received messages from nodes 1 and 2: (2, 2) and (3, 3), respectively. We know that the default aggregation mode is add; and so, would expect an outcome of (5, 5). Indeed, this is the new embedding for node 0.

In a nutshell, thus, the minimal GNN updates every node’s embedding so as to prototypically reflect the node’s neighborhood. Take care though: Nodes represent their neighborhoods, but themselves, they count for nothing. We will change that now.

Adding self loops

All we need to do is modify the adjacency matrix to include edges going from each node back to itself.

from torch_geometric.utils import add_self_loops

class IAmMyOthersAndMyselfAsWell(MessagePassing):
    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes = x.size(0))
        print("in forward, augmented edge index now has shape")
        print(edge_index.shape)
        out = self.propagate(edge_index, x = x)
        return out
    def message(self, x_j):
        return x_j

module = IAmMyOthersAndMyselfAsWell()
out = module(data.x, data.edge_index)
print("result is:")
out

in forward, augmented edge index now has shape:
torch.Size([2, 17])

result is:
tensor([[ 6.,  6.],
        [ 6.,  6.],
        [29., 29.],
        [26., 26.],
        [26., 26.]])

As expected, the neighborhood summary at each node now includes a contribution from each node itself.

Now we know how to access the messages, we’d like to aggregate them in a non-standard way.

Customizing aggregate()

Instead of message(), we now override aggregate(). If we wanted to use another of the “standard” aggregation modes (mean, mul, min, or max), we could just override __init__(), like so:

def __init__(self):
        super().__init__(aggr = "mean")

To implement custom summaries, however, we make use of torch_scatter (one of PyG’s installation prerequisites) for optimal performance. Let me show this by means of a simple example.

from torch_scatter import scatter

class IAmJustTheOppositeReally(MessagePassing):
    def forward(self, x, edge_index):
        out = self.propagate(edge_index, x = x)
        return out
    def aggregate(self, inputs, index):
        print("in aggregate, inputs is")
        # same as x_j (incoming node features)
        print(inputs)
        print("in aggregate, index is")
        # this is data.edge_index[1]
        print(index)
        # see https://pytorch-scatter.readthedocs.io/en/1.3.0/index.html
        # for other aggregation modes
        # default dim is -1
        return - scatter(inputs, index, dim = 0, reduce = "add") 
      
module = IAmJustTheOppositeReally()
out = module(data.x, data.edge_index)
print("result is:")
out

in aggregate, inputs is
tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 1.,  1.],
        [ 3.,  3.],
        [ 2.,  2.],
        [ 3.,  3.],
        [ 3.,  3.],
        [11., 11.],
        [ 3.,  3.],
        [12., 12.],
        [11., 11.],
        [12., 12.]])
        
in aggregate, index is
tensor([1, 0, 2, 0, 2, 1, 3, 2, 4, 2, 4, 3])

result is:
tensor([[ -5.,  -5.],
        [ -4.,  -4.],
        [-26., -26.],
        [-15., -15.],
        [-14., -14.]])

In aggregate(), we have two types of tensors to work with. One, inputs, holds what was returned from message(). In our case, this is identical to x_j, since we didn’t make any modifications to the default behavior. The second, index, holds the recipe for where in the aggregation those features should go. Here, the very first tuple, (1, 1), will contribute to the summary for node 1; the second, (2, 2), to that for node 0 – and so on. By the way, just like x_j (in a single-layer, single-pass setup) is “just” data.x[data.edge_index[0]], that index is “just” data.edge_index[1]. Meaning, this is the list of target nodes connected to the edges in question.

At this point, all kinds of manipulations could be done on either inputs or index; however, we content ourselves with just passing them through to torch_scatter.scatter(), and returning the negated sums. We’ve successfully built a network of contrarians.

By now, we’ve played with message() as well as aggregate(). What about update()?

Adding memory to update()

There’s one thing really strange in what we’re doing. It doesn’t jump to the eye, since we’re not simulating a real training phase; we’ve been calling the layer just once. If we hadn’t, we’d have noticed that at every call, the nodes happily forget who they were before, dutifully assuming the new identities assigned. In reality, we probably want them to evolve in a more consistent way.

For example:

class IDoEvolveOverTime(MessagePassing):
    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        out = self.propagate(edge_index, x = x)
        return out
    def update(self, inputs, x):
        print("in update, inputs is")
        print(inputs)
        print("in update, x is")
        print(x)
        return (inputs + x)/2

module = IDoEvolveOverTime()
out = module(data.x, data.edge_index)
print("result is:")
out

in update, inputs is
tensor([[ 6.,  6.],
        [ 6.,  6.],
        [29., 29.],
        [26., 26.],
        [26., 26.]])
in update, x is
tensor([[ 1.,  1.],
        [ 2.,  2.],
        [ 3.,  3.],
        [11., 11.],
        [12., 12.]])
result is:
tensor([[ 3.5000,  3.5000],
        [ 4.0000,  4.0000],
        [16.0000, 16.0000],
        [18.5000, 18.5000],
        [19.0000, 19.0000]])

In update(), we have access to both the final message aggregate (inputs) and the nodes’ prior states (x). Here, I’m just averaging those two.

At this point, we’ve successfully acquainted ourselves with the three stages of message passing: acting on individual messages, aggregating them, and self-updating based on past state and new information. But none of our models so far could be called a neural network, since there was no learning involved.

Adding parameters

If we look back at the generic message passing formulation:

we see two places where neural network modules can act on the computation: before message aggregation, and as part of the node update process. First, we illustrate the former option. For example, we can apply a MLP in forward(), before the call to aggregate():

from torch.nn import Sequential as Seq, Linear, ReLU

class ILearnAndEvolve(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr = "sum")
        self.mlp = Seq(Linear(in_channels, out_channels),
                       ReLU(),
                       Linear(out_channels, out_channels))
    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        x = self.mlp(x)
        out = self.propagate(edge_index = edge_index, x = x)
        return out
    def update(self, inputs, x):
        return (inputs + x)/2

module = ILearnAndEvolve(2, 2)
out = module(data.x, data.edge_index)
print("result is:")
out

result is:
tensor([[-0.8724, -0.4407],
        [-0.9056, -0.4623],
        [-2.0229, -1.1240],
        [-1.8691, -1.0867],
        [-1.9024, -1.1082]], grad_fn=<DivBackward0>)

Finally, we can apply network modules in both places, as exemplified next.

General message passing

We keep the MLP from the previous class, and add a second in update():

class ILearnAndEvolveDoubly(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr = "sum")
        self.mlp_msg = Seq(Linear(in_channels, out_channels),
                       ReLU(),
                       Linear(out_channels, out_channels))
        self.mlp_upd = Seq(Linear(out_channels, out_channels),
                       ReLU(),
                       Linear(out_channels, out_channels))
    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        x = self.mlp_msg(x)
        out = self.propagate(edge_index = edge_index, x = x)
        return out
    def update(self, inputs, x):
        return self.mlp_upd((inputs + x)/2)

module = ILearnAndEvolveDoubly(2, 2)
out = module(data.x, data.edge_index)
print("result is:")
out

result is:
tensor([[ 0.0573, -0.6988],
        [ 0.0358, -0.6894],
        [-0.1730, -0.6450],
        [-0.5855, -0.4171],
        [-0.5890, -0.4141]], grad_fn=<AddmmBackward0>)

At this point, I hope you’ll feel comfortable to play around, subclassing the MessagePassing base class. Also, if now you consult the above-mentioned documentation page (Creating message passing networks), you’ll be able to map the example implementations (dedicated to popular GNN layer types) to where they “hook into” the message passing process.

Experimentation with MessagePassing was the point of this post. However, you may be wondering: How do I actually use this for node classification? Didn’t the graph have a class defined for each node? (It did: data.y.)

So let me conclude with a (minimal) end-to-end example that uses one of the above modules.

A minimal workflow

To that purpose, we compose that module with a linear one that performs node classification:

class Network(torch.nn.Module):
    def __init__(self, in_channels, out_channels, num_classes):
        super().__init__()
        self.conv = ILearnAndEvolveDoubly(in_channels, out_channels)
        self.classifier = Linear(out_channels, num_classes)
    def forward(self, x, edge_index):
        x, edge_index = data.x, data.edge_index
        x = self.conv(x, edge_index)
        return self.classifier(x)

model = Network(2, 2, 1) 

We can then train the model like any other:

import torch.nn.functional as F

optimizer = torch.optim.Adam(model.parameters(), lr = 0.01)
model.train()

for epoch in range(5):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.binary_cross_entropy_with_logits(out, data.y)
    loss.backward()
    optimizer.step()

preds = torch.sigmoid(out)
preds

tensor([[0.6502],
        [0.6532],
        [0.7027],
        [0.7145],
        [0.7165]], grad_fn=<SigmoidBackward0>)

And that’s it for this time. Stay tuned for examples of how graph models are applied in the sciences, as well as illustrations of bleeding-edge developments in Geometric Deep Learning, the principles-based, heuristics-transcending approach to neural networks.

Thanks for reading!