Introduction to Graphs and Graph Neural Networks

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import seaborn as snss
import pandas as pd
plt.rcParams['xtick.direction'] = 'in'
plt.rcParams['ytick.direction'] = 'in'
plt.rcParams['axes.grid'] = True
plt.rcParams['axes.axisbelow'] = True
plt.rcParams['figure.dpi'] = 150

# Check available hardware
import torch
if torch.cuda.is_available():
    print(f'GPU is available, using GPU {torch.cuda.get_device_name(0)}')
    device = torch.device('cuda')
elif torch.mps.is_available():
    print(f'Apple Metal Performance Shaders framework is available, using {torch.device("mps")}')
    device = torch.device('mps')
else:
    print(f'Hardware acceleration not available, using CPU')
    device = torch.device('cpu')

What is a graph?¶

A graph is a structure containing a set of nodes (also called vertices) and edges that link between nodes. We denote the set of nodes/vertices as $V$ and the edges as $E \subseteq V \times V$ .

A graph $G = (V, E)$ is said to have order $n = |V|$ and size $s = |E|$ . The degree $d$ of a particular vertex is given by the number of edges connected to it.

import networkx as nx

# Define a graph structure using NetworkX
G = nx.Graph()

# A graph has nodes, edges
G.add_nodes_from([
])
print('Nodes:', G.nodes)

G.add_edges_from([
])
print('Edges:', G.edges)

nx.draw(G, with_labels=True, node_color=nx.get_node_attributes(G, 'color').values())

Adjacency matrix¶

The connections in a graph can be defined using an adjacency matrix $A$ . The entries $A_{ij}$ are nonzero if there is an edge between vertices $i$ and $j$ . For a simple undirected graph, the adjacency matrix must be symmetric $A_{ij} = A_{ji}$ .

A = np.asarray(nx.adjacency_matrix(G).todense())
print('Adjacency matrix')
print(A)

Nodes can have values (features)¶

Beyond just specifying connections between elements, graph nodes can contain information encoded in a feature vector $\mathbf{x}$ . Below, see an example where each node contains a scalar feature.

# Nodes can have features/values
G.add_nodes_from([
])
nx.draw(G, labels=nx.get_node_attributes(G, 'feature'), node_size=1e3)

nx.get_node_attributes(G, 'feature').values()

Edges can have values (weights)¶

Likewise, edges in a graph can have features, which are sometimes interpreted as weights on each connection. See example below

# Edges can have features/values
G.add_edges_from([
])
pos = nx.spring_layout(G)
nx.draw(G, pos, labels=nx.get_node_attributes(G, 'feature'), node_size=1e3)
nx.draw_networkx_edge_labels(G,pos, edge_labels=nx.get_edge_attributes(G, 'weight'));

nx.get_edge_attributes(G, 'weight')

Graph edges can be directed¶

For directed graphs (or digraphs), edges are unidirectional

# Graphs can be directed
import networkx as nx

D = nx.DiGraph()

# A graph has nodes, edges
D.add_nodes_from([
])
print('Nodes:', D.nodes)

D.add_edges_from([
])
print('Edges:', D.edges)

nx.draw(D, with_labels=True, node_color=nx.get_node_attributes(D, 'color').values())

# Directed graphs can have unique features on each directed edge
D.add_edges_from([
])
ax = plt.figure().gca()
pos = nx.spring_layout(D)
nx.draw_networkx_nodes(D, pos, node_size=1e3, ax=ax)

connectionstyle=[f"arc3,rad=0.15", f"arc3,rad=0.3"]
nx.draw_networkx_edges(D, pos, arrows=True, arrowstyle='-|>', arrowsize=25, connectionstyle=connectionstyle);
nx.draw_networkx_edge_labels(D,pos, edge_labels=nx.get_edge_attributes(D, 'weight'),
                             connectionstyle=connectionstyle)

plt.gca().axis('off')

In this case the adjancy matrix no longer has to be symmetric!

A = np.asarray(nx.adjacency_matrix(D).todense())
print('Adjacency matrix')
print(A)

Concrete examples¶

link

Graph convolution¶

Let’s start by reviewing some of the key concepts from convolutional neural networks. CNNs leverage parameter sharing and sparse connectivity to achieve translation equivariance. Using local pooling, they can achieve local translation invariance as well. Finally, CNNs learn a hierarchy of features, starting with basic objects like edges and corners in early layers, followed by more complex features in later layers.

Beyond those properties, CNNs are computationally efficient. Their weight sharing means they have a smaller footprint than an FCNN applied to similar data, and the convolution operation can be applied to inputs of arbitrary spatial dimension. That is, a CNN trained on $224 \times 224$ images can be applied to $28 \times 28$ images.

Permutation invariance¶

Graph data does not require translation invariance, because there is no notion of spatial position. The only things graphs have are node features, adjacency matrix, and optionally edge features. They also require permutation invariance - if the order of the nodes are scrambled, the output of an operation must not change.

# Permutation invariance means that an operation should give the same outputs for the following three graphs
import networkx as nx

G1 = nx.Graph()

# A graph has nodes, edges
G1.add_nodes_from([
])
print('Nodes:', G1.nodes)

G1.add_edges_from([
])
print('Edges:', G1.edges)
print('Adjacency Matrix:\n', np.asarray(nx.adjacency_matrix(G1).todense()))

# Permutation invariance means that an operation should give the same outputs for the following three graphs
import networkx as nx

G2 = nx.Graph()

# A graph has nodes, edges
G2.add_nodes_from([
])
print('Nodes:', G2.nodes)

G2.add_edges_from([
])
print('Edges:', G2.edges)
print('Adjacency Matrix:\n', np.asarray(nx.adjacency_matrix(G2).todense()))

# Permutation invariance means that an operation should give the same outputs for the following three graphs
import networkx as nx

G3 = nx.Graph()

# A graph has nodes, edges
G3.add_nodes_from([
])
print('Nodes:', G3.nodes)

G3.add_edges_from([
])
print('Edges:', G3.edges)
print('Adjacency Matrix:\n', np.asarray(nx.adjacency_matrix(G3).todense()))

nx.draw(G1, with_labels=True, node_color=nx.get_node_attributes(G1, 'color').values())

nx.draw(G2, with_labels=True, node_color=nx.get_node_attributes(G2, 'color').values())

nx.draw(G3, with_labels=True, node_color=nx.get_node_attributes(G3, 'color').values())

feature_map = {
    'red': 0,
    'green': 1,
    'blue': 2,
    'orange': 3
}

def get_one_hot_features(G):
    features = np.array([feature_map[v] for v in nx.get_node_attributes(G, 'color').values()])
    one_hot = np.zeros([features.shape[0], len(np.unique(features))])
    one_hot[np.arange(len(one_hot)), features] = 1.
    return one_hot

print(get_one_hot_features(G1))
print(get_one_hot_features(G2))
print(get_one_hot_features(G3))

def print_nodes_features(nodes, features):
    for node, feat in zip(nodes, features):
        print(node, ':', feat)

print_nodes_features(G1.nodes(), get_one_hot_features(G1))
print()
print_nodes_features(G2.nodes(), get_one_hot_features(G2))
print()
print_nodes_features(G3.nodes(), get_one_hot_features(G3))

A basic graph convolution¶

One option is to sum over the incoming features at each node to determine the updated feature value. We can apply a weight matrix $W_{ij}$ to each incoming feature, but it will be the same weight matrix. For example, the preactivation $\mathbf{z}_i$ at node $i$ will be

\begin{align*} \mathbf{z}_i &= \mathbf{b} + \mathbf{x}_i \, \mathbf{W}_1+ \sum_{j \in \mathcal{N}(i)} \mathbf{x}_j \, \mathbf{W}_2 \end{align*}

Here, $\mathbf{b}$ is a learnable bias, $\mathbf{W}_1$ is a matrix applied to the local feature, and $\mathbf{W}_2$ is a matrix applied to each neighbor node’s feature. In principle, $\mathbf{W}_1,\mathbf{W}_2$ can be matrices of shape $f_{\text{in}} \times f_{\text{out}}$ to change the feature dimension at each update.

The term $j \in \mathcal{N}(i)$ is a sum over the neighbors of node $i$ . This could be computed using the adjacency matrix, i.e.

\begin{align*} \sum_{j \in \mathcal{N}(i)} \mathbf{x}_j = \sum_{j} A_{ij} \mathbf{x}_j \end{align*}

np.random.seed(42)
f_in, f_out = 4, 4
W1 = np.random.randint(low=0, high=5, size=[f_in, f_out])
W2 = np.random.randint(low=0, high=5, size=[f_in, f_out])
b = np.random.randint(low=0, high=2, size=[f_out])

G = G1
X = get_one_hot_features(G)
A = nx.adjacency_matrix(G)

z = 
print_nodes_features(G.nodes(), z)

G = G2
X = get_one_hot_features(G)
A = nx.adjacency_matrix(G)

z = 
print_nodes_features(G.nodes(), z)

G = G3
X = get_one_hot_features(G)
A = nx.adjacency_matrix(G)

z = 
print_nodes_features(G.nodes(), z)

More general graph convolutions¶

There are many ways to perform a graph convolution. The version above was the form proposed by Kipf and Welling in 2016 (link). A more general way to view this is as a two-step process:

Message passing - each node gathers information from its neighbors
Node update - each node updates its features based on the incoming messages

The message passing phase can be framed mathematically as

\begin{align*} \mathbf{m}_i = \sum_{j \in \mathcal{N}(i)} M \left( \mathbf{x}_i, \mathbf{x}_j, e_{ij} \right) \end{align*}

Here, $M$ is the message function that determines the message from each neighbor $j$ based on their hidden states, and optionally any features or weights on the incoming edge. If there are no edge features, then we use $e_{ij} = A_{ij}$ .

In our previous example, the message function was simply

\begin{align*} M \left( \mathbf{x}_i, \mathbf{x}_j, e_{ij} \right) = \mathbf{x}_j \mathbf{W}_2 \end{align*}

Other flavors of message functions include attention mechanisms

\begin{align*} M \left( \mathbf{x}_i, \mathbf{x}_j, e_{ij} \right) &= a(\mathbf{x}_i, \mathbf{x}_j) \, \psi (\mathbf{x}_j) \end{align*}

where $a(\mathbf{x}_i, \mathbf{x}_j)$ is a learnable function that computes the importance of an incoming value $\psi(\mathbf{x}_j)$ . A common formulation is scaled-dot product attention

\begin{align*} M \left( \mathbf{x}_i, \mathbf{x}_j, e_{ij} \right) &= \text{softmax} \left( \frac{\left( \mathbf{x}_i \mathbf{W}_Q \right) \cdot \left( \mathbf{x}_j \mathbf{W}_K \right)}{\sqrt{f_k}} \right) \left( \mathbf{x}_j \mathbf{W}_V \right) \end{align*}

Here, $\mathbf{W}_Q, \mathbf{W}_K \in \mathbb{R}^{f_{\text{in}} \times f_{k}}$ are query and key weight matrices and $\mathbf{W}_V \in \mathbb{R}^{f_{\text{in}} \times f_{\text{out}}}$ is a value weight matrix.

The node update phase uses an update function to compute the next feature value at the particular node.

\begin{align*} \mathbf{x}^{(t+1)}_i = U_t \left( \mathbf{m}^{(t)}_i, \mathbf{x}^{(t)}_i \right) \end{align*}

In our previous example, the node update function was simply

\begin{align*} \mathbf{x}^{(t+1)}_i = \mathbf{x}^{(t)}_i \mathbf{W}_1 + \mathbf{m}_i + \mathbf{b} \end{align*}

import torch
import torch.nn as nn
import torch.nn.functional as F

class GraphConvolutionLayer(nn.Module):
    def __init__(self, f_in, f_out):
        super().__init__()

    def forward(self, X, A):

layer = GraphConvolutionLayer(4, 8)

G = G3
X = torch.FloatTensor(get_one_hot_features(G))
A = torch.FloatTensor(nx.adjacency_matrix(G).todense())

with torch.no_grad():
    Xt = layer(X, A)

print_nodes_features(G.nodes(), Xt)

G = G2
X = torch.FloatTensor(get_one_hot_features(G))
A = torch.FloatTensor(nx.adjacency_matrix(G).todense())

with torch.no_grad():
    Xt = layer(X, A)

print_nodes_features(G.nodes(), Xt)

# Stacking graph convolutional layers
# This increases the "receptive field" of each node
gcn = nn.ModuleList([
])

with torch.no_grad():
    Xt = X.clone()
    for layer in gcn:
        Xt = layer(Xt, A)
        print_nodes_features(G.nodes(), Xt)
        print()

Graph pooling¶

Graph convolution, in its many forms, does not change the structure of the graph. That is, there are always the same number of nodes and edges. Instead, graph convolutional layers update the features within the graph through repeated computations.

It may be useful to build analogies to local max pooling in CNNs. We may talk about that in later lectures, but this is in general a complicated process. Instead, we’ll use global sum, average, or max pooling which aggregates the feature values over the entire graph

class GlobalSumPool(nn.Module):
    def forward(self, X, batch_matrix=None):
        """ Go to the board for this
            batch_matrix[i,j] = 1 if node j belongs to graph i
            batch_matrix[i,j] = 0 if node j does not belong to graph i
            Useful for batching multiple graphs of different sizes
        """

pool = GlobalSumPool()
pool(Xt)

Putting it all together¶

# A Graph neural network that might be sufficient for classifying based on graph data
class GraphClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.gcn = 
        self.pool = 
        self.fc = 

    def forward(self, X, A, batch_matrix=None):

model = GraphClassifier()
model(X, A)

Building a batch of data¶

for G in [G1, G2, G3]:
    print(G.nodes(), nx.adjacency_matrix(G))
    print()

class GraphDataset(torch.utils.data.Dataset):
    def __init__(self, graphs):
        self.graphs = graphs

    def __len__(self):
        return len(self.graphs)

    def __getitem__(self, idx):
        return self.graphs[idx]

dataset = GraphDataset([G1, G2, G3])

loader = torch.utils.data.DataLoader(dataset, batch_size=3, shuffle=False)

batch = next(iter(loader))

def collate_graphs(batch):
    """ Return a batch of node features, adjacency matrices, and the batch_matrix
    """
    X = []
    A = []
    sizes = []
    
loader = torch.utils.data.DataLoader(dataset, batch_size=3, shuffle=False, collate_fn=collate_graphs)
Xb, Ab, Bb = next(iter(loader))

print('Batch features')
print(Xb)
print()
print('Batch adjacency')
print(Ab)
print()
print('Batch matrix')
print(Bb)
print()

model = GraphClassifier()
model(Xb, Ab, Bb)