Graph Neural Networks: обучение на графовых данных
opensourceaimlgnndeep-learninggraphsit
Введение: Графы в машинном обучении
Почему графы?
Большинство ML работает с:
- Таблицы (CSV, SQL)
- Изображения (сетка пикселей)
- Последовательности (текст, время)
Реальность - графы:
- Социальные сети (люди + связи)
- Молекулы (атомы + связи)
- Рекомендации (пользователи + товары)
- Дороги (перекрёстки + дороги)
- Знания (сущности + отношения)
Факт: 70% данных в enterprise - это графовые данные.
Основы графов
Определение графа
Graph G = (V, E)
V = {v1, v2, ..., vn} # Вершины (nodes)
E = {(v1,v2), (v2,v3), ...} # Рёбра (edges)
Пример: социальная сеть
V = {Alice, Bob, Charlie, Diana}
E = {(Alice,Bob), (Bob,Charlie), (Charlie,Diana), (Alice,Diana)}
Свойства:
- Направленный / ненаправленный
- Взвешенный / невзвешенный
- Атрибуты вершин и рёбер
Матричное представление
import torch
import torch_geometric
# Adjacency matrix
# A B C D
A = [[0, 1, 0, 1], # Alice
[1, 0, 1, 0], # Bob
[0, 1, 0, 1], # Charlie
[1, 0, 1, 0]] # Diana
# Feature matrix
# feature1 feature2
X = [[1, 0], # Alice
[0, 1], # Bob
[1, 1], # Charlie
[0, 0]] # Diana
# In PyTorch Geometric:
from torch_geometric.data import Data
edge_index = torch.tensor([
[0, 1, 1, 2, 2, 3, 3, 0], # source
[1, 0, 2, 1, 3, 2, 0, 3] # target
])
x = torch.tensor([[1., 0.],
[0., 1.],
[1., 1.],
[0., 0.]], dtype=torch.float)
data = Data(x=x, edge_index=edge_index)
Graph Neural Networks
Message Passing
Core GNN operation: Message Passing
For each edge (j, i):
1. Message: m_ji = φ(h_j, h_i, e_ji)
2. Aggregate: m_i = Σ_j m_ji
3. Update: h_i' = ψ(h_i, m_i)
Where:
h = node features
e = edge features
φ, ψ = learnable functions
class MessagePassing(torch.nn.Module):
def __init__(self, hidden_dim):
super().__init__()
self.message_fn = torch.nn.Sequential(
torch.nn.Linear(hidden_dim * 2 + edge_dim, hidden_dim),
torch.nn.ReLU(),
)
self.update_fn = torch.nn.GRUCell(hidden_dim, hidden_dim)
def forward(self, x, edge_index, edge_attr):
"""
x: (num_nodes, hidden_dim)
edge_index: (2, num_edges)
edge_attr: (num_edges, edge_dim)
"""
src, dst = edge_index # source and target nodes
# Create messages for each edge
x_src = x[src] # (num_edges, hidden_dim)
x_dst = x[dst] # (num_edges, hidden_dim)
# Combine features
combined = torch.cat([x_src, x_dst, edge_attr], dim=-1)
messages = self.message_fn(combined) # (num_edges, hidden_dim)
# Aggregate messages (sum)
out = torch.zeros_like(x)
out.index_add_(0, dst, messages)
# Update node states
h_new = self.update_fn(out, x)
return h_new
Graph Convolution (GCN)
GCN: Spectral graph convolution (approximation)
Update rule:
H' = σ(D^(-1/2) A D^(-1/2) H W)
Where:
A = adjacency matrix (with self-loops)
D = degree matrix
H = node features
W = learnable weights
σ = activation function
Simplified:
h_i' = σ(Σ_j∈N[i] (1/c_ij) h_j W)
N[i] = neighbors of i (including self)
c_ij = normalization constant
class GCNLayer(torch.nn.Module):
def __init__(self, in_features, out_features):
super().__init__()
self.linear = torch.nn.Linear(in_features, out_features)
self.activation = torch.nn.ReLU()
def forward(self, x, edge_index, edge_weight=None):
"""
x: (num_nodes, in_features)
edge_index: (2, num_edges)
edge_weight: (num_edges,) optional
"""
# Prepare adjacency
row, col = edge_index
num_nodes = x.shape[0]
# Add self-loops
self_loop = torch.arange(num_nodes).to(edge_index)
edge_index_full = torch.cat([
torch.stack([self_loop, self_loop], dim=0),
edge_index
], dim=1)
# Compute normalization
deg = degree(edge_index_full[1], num_nodes)
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
# Normalized adjacency
for i in range(edge_index_full.shape[1]):
edge_weight_norm[i] = deg_inv_sqrt[edge_index_full[0, i]] * \
deg_inv_sqrt[edge_index_full[1, i]]
# Message passing
out = self.linear(x)
# Aggregate
row, col = edge_index_full
out_col = out[col]
out_col = out_col * edge_weight_norm.unsqueeze(1)
agg = torch.zeros_like(out)
agg.index_add_(0, row, out_col)
return self.activation(agg)
class GCN(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels):
super().__init__()
self.conv1 = GCNLayer(in_channels, hidden_channels)
self.conv2 = GCNLayer(hidden_channels, out_channels)
def forward(self, x, edge_index):
x = self.conv1(x, edge_index)
x = torch.nn.functional.relu(x)
x = torch.nn.functional.dropout(x, training=self.training)
x = self.conv2(x, edge_index)
return x
GraphSAGE
GraphSAGE: Inductive representation learning
Key idea: Learn aggregation functions
instead of fixed graph Laplacian
Aggregators:
- Mean (original)
- LSTM
- Pooling
Works on unseen nodes!
class SAGELayer(torch.nn.Module):
def __init__(self, in_features, out_features, aggregator='mean'):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.aggregator = aggregator
# Self features transformation
self.self_linear = torch.nn.Linear(in_features, out_features)
# Aggregation
if aggregator == 'mean':
self.aggregator_linear = torch.nn.Linear(in_features, out_features)
elif aggregator == 'pool':
self.pool = SAGEPool(in_features, out_features)
elif aggregator == 'lstm':
self.lstm = nn.LSTM(in_features, out_features, batch_first=True)
self.activation = torch.nn.ReLU()
def forward(self, x, edge_index):
row, col = edge_index
num_nodes = x.shape[0]
# Transform self
h_self = self.self_linear(x)
# Aggregate neighbors
h_neighbors = self.aggregate(x, row, col)
# Combine
out = torch.cat([h_self, h_neighbors], dim=-1)
out = self.activation(out)
# L2 normalize (like original paper)
out = torch.nn.functional.normalize(out, p=2, dim=1)
return out
def aggregate(self, x, row, col):
if self.aggregator == 'mean':
# Mean aggregator
neighbor_features = x[col]
deg = torch.bincount(row, minlength=x.shape[0])
deg_inv = 1.0 / deg.clamp(min=1)
out = torch.zeros_like(x)
out.index_add_(0, row, neighbor_features * deg_inv[row].unsqueeze(1))
return self.aggregator_linear(out)
elif self.aggregator == 'pool':
return self.pool(x, row, col)
GAT (Graph Attention)
GAT: Self-attention on graphs
For each node i:
h_i' = Σ_j α_ij Wh_j
α_ij = softmax_i(e_ij)
e_ij = LeakyReLU(a^T [Wh_i || Wh_j])
Where:
|| = concatenation
a = learnable attention vector
class GATLayer(torch.nn.Module):
def __init__(self, in_features, out_features, n_heads=1, dropout=0.6):
super().__init__()
self.n_heads = n_heads
self.head_dim = out_features // n_heads
# Linear projection
self.W = torch.nn.Linear(in_features, out_features)
# Attention mechanism
self.a = torch.nn.Parameter(
torch.randn(2 * self.head_dim, n_heads)
)
self.leakyrelu = torch.nn.LeakyReLU(0.2)
self.softmax = torch.nn.Softmax(dim=-1)
self.dropout = torch.nn.Dropout(dropout)
def forward(self, x, edge_index):
"""
x: (num_nodes, in_features)
edge_index: (2, num_edges)
"""
num_nodes = x.shape[0]
# Linear transform
H = self.W(x) # (num_nodes, out_features)
H = H.view(num_nodes, self.n_heads, self.head_dim)
# Compute attention coefficients
row, col = edge_index
# Attention for source and target
a_src = (H[row] * self.a[:self.head_dim]).sum(dim=-1) # (num_edges, n_heads)
a_dst = (H[col] * self.a[self.head_dim:]).sum(dim=-1) # (num_edges, n_heads)
e = self.leakyrelu(a_src + a_dst.T) # (num_edges, n_heads)
# Sparse attention
alpha = torch.zeros(num_nodes, num_nodes, self.n_heads, device=x.device)
alpha[row, col] = self.dropout(self.softmax(e))
# Aggregate
out = torch.zeros(num_nodes, self.n_heads * self.head_dim, device=x.device)
for h in range(self.n_heads):
H_h = H[:, h, :] # (num_nodes, head_dim)
alpha_h = alpha[:, :, h] # (num_nodes, num_nodes)
out[:, h*self.head_dim:(h+1)*self.head_dim] = (alpha_h @ H_h)
return out
GNN Training
Node Classification
class GNNClassifier(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, num_classes):
super().__init__()
self.encoder = GCN(in_channels, hidden_channels, hidden_channels)
self.classifier = torch.nn.Linear(hidden_channels, num_classes)
def forward(self, x, edge_index):
embeddings = self.encoder(x, edge_index)
logits = self.classifier(embeddings)
return logits
# Training loop
model = GNNClassifier(in_channels=128, hidden_channels=256, num_classes=7)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss()
for epoch in range(200):
model.train()
logits = model(data.x, data.edge_index)
loss = criterion(logits[train_mask], data.y[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Evaluate
model.eval()
with torch.no_grad():
logits = model(data.x, data.edge_index)
pred = logits.argmax(dim=1)
accuracy = (pred[test_mask] == data.y[test_mask]).float().mean()
Graph Classification
class GraphClassifier(torch.nn.Module):
def __init__(self, in_channels, hidden_channels, num_classes):
super().__init__()
self.gnn = GCN(in_channels, hidden_channels, hidden_channels)
# Readout functions
self.readout = 'mean' # mean, sum, max, attention
self.classifier = torch.nn.Sequential(
torch.nn.Linear(hidden_channels, hidden_channels),
torch.nn.ReLU(),
torch.nn.Dropout(0.5),
torch.nn.Linear(hidden_channels, num_classes)
)
def forward(self, batch_data):
"""
batch_data: batch of graphs
"""
# GNN encoding
x = torch.cat([data.x for data in batch_data])
edge_index = torch.cat([data.edge_index for data in batch_data], dim=1)
batch = torch.cat([torch.full((data.x.shape[0],), i)
for i, data in enumerate(batch_data)])
h = self.gnn(x, edge_index)
# Readout (pooling)
if self.readout == 'mean':
graph_repr = scatter(h, batch, dim=0, reduce='mean')
elif self.readout == 'sum':
graph_repr = scatter(h, batch, dim=0, reduce='sum')
elif self.readout == 'max':
graph_repr = scatter(h, batch, dim=0, reduce='max')
return self.classifier(graph_repr)
Applications
Molecular Property Prediction
Molecule = Graph
Nodes = Atoms (element, charge, hybridization)
Edges = Bonds (type, length, angle)
Tasks:
- Drug discovery (toxicity, binding affinity)
- Material science (conductivity, stability)
- Chemistry (reaction prediction)
class MoleculeGNN(torch.nn.Module):
def __init__(self, num_atom_types, num_bond_types, hidden_dim):
super().__init__()
# Atom embeddings
self.atom_embedding = torch.nn.Embedding(num_atom_types, hidden_dim)
# Bond embeddings
self.bond_embedding = torch.nn.Embedding(num_bond_types, hidden_dim)
# GNN layers
self.gnn_layers = torch.nn.ModuleList([
GATLayer(hidden_dim, hidden_dim)
for _ in range(3)
])
# Property prediction heads
self.property_heads = torch.nn.ModuleDict({
'solubility': torch.nn.Linear(hidden_dim, 1),
'toxicity': torch.nn.Linear(hidden_dim, 1),
'binding': torch.nn.Linear(hidden_dim, 1),
})
def forward(self, molecule_data):
# Node features
atom_types = molecule_data.x.long()
atom_features = self.atom_embedding(atom_types)
# Edge features
bond_types = molecule_data.edge_attr.long()
bond_features = self.bond_embedding(bond_types)
# Message passing
h = atom_features
for gnn_layer in self.gnn_layers:
h = gnn_layer(h, molecule_data.edge_index)
h = torch.nn.functional.relu(h)
# Graph-level readout
batch = molecule_data.batch
graph_repr = scatter(h, batch, dim=0, reduce='mean')
# Predict properties
properties = {
name: head(graph_repr)
for name, head in self.property_heads.items()
}
return properties
Recommendation Systems
User-Item Graph:
Users: age, gender, location
Items: category, price, features
Edges: purchases, clicks, ratings
GNN approach:
1. Encode user and item features
2. Message passing on bipartite graph
3. Predict edge probability (will user click item?)
GNN Frameworks
# PyTorch Geometric
import torch_geometric.nn as pyg_nn
# Common layers
gcnn = pyg_nn.GCNConv(in_channels, out_channels)
sage = pyg_nn.SAGEConv(in_channels, out_channels)
gat = pyg_nn.GATv2Conv(in_channels, out_channels)
gin = pyg_nn.GINConv(torch.nn.Sequential(
torch.nn.Linear(in_channels, out_channels),
torch.nn.ReLU(),
torch.nn.Linear(out_channels, out_channels)
))
# Pooling layers
mean_pool = pyg_nn.global_mean_pool
sum_pool = pyg_nn.global_add_pool
max_pool = pyg_nn.global_max_pool
attention_pool = pyg_nn.GlobalAttention(torch.nn.LSTM(1, 1))
# Datasets
from torch_geometric.datasets import ZINC, OGBNaMolecules, Planetoid
from torch_geometric.loader import DataLoader
dataset = Planetoid(root='data/Cora', name='Cora')
loader = DataLoader(dataset, batch_size=32, shuffle=True)
Заключение
GNNs provide:
- Relational inductive bias: Learn from structure
- Message passing: Information flows through graph
- Flexibility: Works on any graph-structured data
Key architectures:
- GCN - spectral convolution
- GraphSAGE - inductive aggregation
- GAT - attention-based
- GIN - theoretically powerful
Applications:
- Drug discovery
- Recommendation systems
- Social network analysis
- Knowledge graphs
- Materials science