mirrors/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/core/Graphormer/README.md
+33
View File
@@ -0,0 +1,33 @@
Graphormer
==============================
## Introduction
* Graphormer is a Transformer model designed for graph-structured data, which encodes the structural information of a graph into the standard Transformer. Specifically, Graphormer utilizes Degree Encoding to measure the importance of nodes, Spatial Encoding and Path Encoding to measure the relation between node pairs. The former plus the node features serve as input to Graphormer, while the latter acts as bias terms in the self-attention module.
* paper link: [https://arxiv.org/abs/2106.05234](https://arxiv.org/abs/2106.05234)
## Requirements
- accelerate
- transformers
- ogb
## Dataset
Task: Graph Property Prediction
| Dataset | #Graphs | #Node Feats | #Edge Feats | Metric |
| :---------: | :-----: | :---------: | :---------: | :-----: |
| ogbg-molhiv | 41,127 | 9 | 3 | ROC-AUC |
How to run
----------
```bash
accelerate launch --multi_gpu --mixed_precision=fp16 main.py
```
> **_NOTE:_** The script will automatically download weights pre-trained on PCQM4Mv2. To reproduce the same result, set the total batch size to 64.
## Summary
* ogbg-molhiv (pretrained on PCQM4Mv2): ~0.791
examples/core/Graphormer/dataset.py
+114
View File
@@ -0,0 +1,114 @@
"""
This file contains the MolHIVDataset class, which handles data preprocessing
(computing required graph features, converting graphs to tensors) of the
ogbg-molhiv dataset.
"""
import torch as th
import torch.nn.functional as F
from dgl import shortest_dist
from ogb.graphproppred import DglGraphPropPredDataset
from torch.nn.utils.rnn import pad_sequence
class MolHIVDataset(th.utils.data.Dataset):
def __init__(self):
dataset = DglGraphPropPredDataset(name="ogbg-molhiv")
split_idx = dataset.get_idx_split()
# Compute the shortest path distances and their corresponding paths
# of all graphs during preprocessing.
for g, label in dataset:
spd, path = shortest_dist(g, root=None, return_paths=True)
g.ndata["spd"] = spd
g.ndata["path"] = path
self.train, self.val, self.test = (
dataset[split_idx["train"]],
dataset[split_idx["valid"]],
dataset[split_idx["test"]],
)
def collate(self, samples):
# To match Graphormer's input style, all graph features should be
# padded to the same size. Keep in mind that different graphs may
# have varying feature sizes since they have different number of
# nodes, so they will be aligned with the graph having the maximum
# number of nodes.
graphs, labels = map(list, zip(*samples))
labels = th.stack(labels)
num_graphs = len(graphs)
num_nodes = [g.num_nodes() for g in graphs]
max_num_nodes = max(num_nodes)
# Graphormer adds a virual node to the graph, which is connected to
# all other nodes and supposed to represent the graph embedding. So
# here +1 is for the virtual node.
attn_mask = th.zeros(num_graphs, max_num_nodes + 1, max_num_nodes + 1)
node_feat = []
in_degree, out_degree = [], []
path_data = []
# Since shortest_dist returns -1 for unreachable node pairs and padded
# nodes are unreachable to others, distance relevant to padded nodes
# use -1 padding as well.
dist = -th.ones(
(num_graphs, max_num_nodes, max_num_nodes), dtype=th.long
)
for i in range(num_graphs):
# A binary mask where invalid positions are indicated by True.
attn_mask[i, :, num_nodes[i] + 1 :] = 1
# +1 to distinguish padded non-existing nodes from real nodes
node_feat.append(graphs[i].ndata["feat"] + 1)
in_degree.append(
th.clamp(graphs[i].in_degrees() + 1, min=0, max=512)
)
out_degree.append(
th.clamp(graphs[i].out_degrees() + 1, min=0, max=512)
)
# Path padding to make all paths to the same length "max_len".
path = graphs[i].ndata["path"]
path_len = path.size(dim=2)
# shape of shortest_path: [n, n, max_len]
max_len = 5
if path_len >= max_len:
shortest_path = path[:, :, :max_len]
else:
p1d = (0, max_len - path_len)
# Use the same -1 padding as shortest_dist for
# invalid edge IDs.
shortest_path = F.pad(path, p1d, "constant", -1)
pad_num_nodes = max_num_nodes - num_nodes[i]
p3d = (0, 0, 0, pad_num_nodes, 0, pad_num_nodes)
shortest_path = F.pad(shortest_path, p3d, "constant", -1)
# +1 to distinguish padded non-existing edges from real edges
edata = graphs[i].edata["feat"] + 1
# shortest_dist pads non-existing edges (at the end of shortest
# paths) with edge IDs -1, and th.zeros(1, edata.shape[1]) stands
# for all padded edge features.
edata = th.cat(
(edata, th.zeros(1, edata.shape[1]).to(edata.device)), dim=0
)
path_data.append(edata[shortest_path])
dist[i, : num_nodes[i], : num_nodes[i]] = graphs[i].ndata["spd"]
# node feat padding
node_feat = pad_sequence(node_feat, batch_first=True)
# degree padding
in_degree = pad_sequence(in_degree, batch_first=True)
out_degree = pad_sequence(out_degree, batch_first=True)
return (
labels.reshape(num_graphs, -1),
attn_mask,
node_feat,
in_degree,
out_degree,
th.stack(path_data),
dist,
)
examples/core/Graphormer/main.py
+274
View File
@@ -0,0 +1,274 @@
"""
This script finetunes and tests a Graphormer model (pretrained on PCQM4Mv2)
for graph classification on ogbg-molhiv dataset.
Paper: [Do Transformers Really Perform Bad for Graph Representation?]
(https://arxiv.org/abs/2106.05234)
This flowchart describes the main functional sequence of the provided example.
main
│
└───> train_val_pipeline
│
├───> Load and preprocess dataset
│
├───> Download pretrained model
│
├───> train_epoch
│ │
│ └───> Graphormer.forward
│
└───> evaluate_network
│
└───> Graphormer.inference
"""
import argparse
import random
import torch as th
import torch.nn as nn
from accelerate import Accelerator
from dataset import MolHIVDataset
from dgl.data import download
from dgl.dataloading import GraphDataLoader
from model import Graphormer
from ogb.graphproppred import Evaluator
from transformers.optimization import (
AdamW,
get_polynomial_decay_schedule_with_warmup,
)
# Instantiate an accelerator object to support distributed
# training and inference.
accelerator = Accelerator()
def train_epoch(model, optimizer, data_loader, lr_scheduler):
model.train()
epoch_loss = 0
list_scores = []
list_labels = []
loss_fn = nn.BCEWithLogitsLoss()
for (
batch_labels,
attn_mask,
node_feat,
in_degree,
out_degree,
path_data,
dist,
) in data_loader:
optimizer.zero_grad()
device = accelerator.device
batch_scores = model(
node_feat.to(device),
in_degree.to(device),
out_degree.to(device),
path_data.to(device),
dist.to(device),
attn_mask=attn_mask,
)
loss = loss_fn(batch_scores, batch_labels.float())
accelerator.backward(loss)
optimizer.step()
lr_scheduler.step()
epoch_loss += loss.item()
list_scores.append(batch_scores)
list_labels.append(batch_labels)
# Release GPU memory.
del (
batch_labels,
batch_scores,
loss,
attn_mask,
node_feat,
in_degree,
out_degree,
path_data,
dist,
)
th.cuda.empty_cache()
epoch_loss /= len(data_loader)
evaluator = Evaluator(name="ogbg-molhiv")
epoch_auc = evaluator.eval(
{"y_pred": th.cat(list_scores), "y_true": th.cat(list_labels)}
)["rocauc"]
return epoch_loss, epoch_auc
def evaluate_network(model, data_loader):
model.eval()
epoch_loss = 0
loss_fn = nn.BCEWithLogitsLoss()
with th.no_grad():
list_scores = []
list_labels = []
for (
batch_labels,
attn_mask,
node_feat,
in_degree,
out_degree,
path_data,
dist,
) in data_loader:
device = accelerator.device
batch_scores = model(
node_feat.to(device),
in_degree.to(device),
out_degree.to(device),
path_data.to(device),
dist.to(device),
attn_mask=attn_mask,
)
# Gather all predictions and targets.
all_predictions, all_targets = accelerator.gather_for_metrics(
(batch_scores, batch_labels)
)
loss = loss_fn(all_predictions, all_targets.float())
epoch_loss += loss.item()
list_scores.append(all_predictions)
list_labels.append(all_targets)
epoch_loss /= len(data_loader)
evaluator = Evaluator(name="ogbg-molhiv")
epoch_auc = evaluator.eval(
{"y_pred": th.cat(list_scores), "y_true": th.cat(list_labels)}
)["rocauc"]
return epoch_loss, epoch_auc
def train_val_pipeline(params):
dataset = MolHIVDataset()
accelerator.print(
f"train, test, val sizes: {len(dataset.train)}, "
f"{len(dataset.test)}, {len(dataset.val)}."
)
accelerator.print("Finished loading.")
train_loader = GraphDataLoader(
dataset.train,
batch_size=params.batch_size,
shuffle=True,
collate_fn=dataset.collate,
pin_memory=True,
num_workers=16,
)
val_loader = GraphDataLoader(
dataset.val,
batch_size=params.batch_size,
shuffle=False,
collate_fn=dataset.collate,
pin_memory=True,
num_workers=16,
)
test_loader = GraphDataLoader(
dataset.test,
batch_size=params.batch_size,
shuffle=False,
collate_fn=dataset.collate,
pin_memory=True,
num_workers=16,
)
# Load pre-trained model.
download(url="https://data.dgl.ai/pre_trained/graphormer_pcqm.pth")
model = Graphormer()
state_dict = th.load("graphormer_pcqm.pth")
model.load_state_dict(state_dict)
model.reset_output_layer_parameters()
num_epochs = 16
total_updates = 33000 * num_epochs / params.batch_size
# Use warmup schedule to avoid overfitting at the very beginning
# of training, the ratio 0.16 is the same as the paper.
warmup_updates = total_updates * 0.16
optimizer = AdamW(model.parameters(), lr=1e-4, eps=1e-8, weight_decay=0)
lr_scheduler = get_polynomial_decay_schedule_with_warmup(
optimizer,
num_warmup_steps=warmup_updates,
num_training_steps=total_updates,
lr_end=1e-9,
power=1.0,
)
epoch_train_AUCs, epoch_val_AUCs, epoch_test_AUCs = [], [], []
# Pass all objects relevant to training to the prepare() method as required
# by Accelerate.
(
model,
optimizer,
train_loader,
val_loader,
test_loader,
lr_scheduler,
) = accelerator.prepare(
model, optimizer, train_loader, val_loader, test_loader, lr_scheduler
)
for epoch in range(num_epochs):
epoch_train_loss, epoch_train_auc = train_epoch(
model, optimizer, train_loader, lr_scheduler
)
epoch_val_loss, epoch_val_auc = evaluate_network(model, val_loader)
epoch_test_loss, epoch_test_auc = evaluate_network(model, test_loader)
epoch_train_AUCs.append(epoch_train_auc)
epoch_val_AUCs.append(epoch_val_auc)
epoch_test_AUCs.append(epoch_test_auc)
accelerator.print(
f"Epoch={epoch + 1} | train_AUC={epoch_train_auc:.3f} | "
f"val_AUC={epoch_val_auc:.3f} | test_AUC={epoch_test_auc:.3f}"
)
# Return test and train AUCs with best val AUC.
index = epoch_val_AUCs.index(max(epoch_val_AUCs))
val_auc = epoch_val_AUCs[index]
train_auc = epoch_train_AUCs[index]
test_auc = epoch_test_AUCs[index]
accelerator.print("Test ROCAUC: {:.4f}".format(test_auc))
accelerator.print("Val ROCAUC: {:.4f}".format(val_auc))
accelerator.print("Train ROCAUC: {:.4f}".format(train_auc))
accelerator.print("Best epoch index: {:.4f}".format(index))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--seed",
default=1,
type=int,
help="Please give a value for random seed",
)
parser.add_argument(
"--batch_size",
default=16,
type=int,
help="Please give a value for batch_size",
)
args = parser.parse_args()
# Set manual seed to bind the order of training data to the random seed.
random.seed(args.seed)
th.manual_seed(args.seed)
if th.cuda.is_available():
th.cuda.manual_seed(args.seed)
train_val_pipeline(args)
examples/core/Graphormer/model.py
+139
View File
@@ -0,0 +1,139 @@
"""
This file defines the Graphormer model, which utilizes DegreeEncoder,
SpatialEncoder, PathEncoder and GraphormerLayer from DGL build-in modules.
"""
import torch as th
import torch.nn as nn
from dgl.nn import DegreeEncoder, GraphormerLayer, PathEncoder, SpatialEncoder
class Graphormer(nn.Module):
def __init__(
self,
num_classes=1,
edge_dim=3,
num_atoms=4608,
max_degree=512,
num_spatial=511,
multi_hop_max_dist=5,
num_encoder_layers=12,
embedding_dim=768,
ffn_embedding_dim=768,
num_attention_heads=32,
dropout=0.1,
pre_layernorm=True,
activation_fn=nn.GELU(),
):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
self.embedding_dim = embedding_dim
self.num_heads = num_attention_heads
self.atom_encoder = nn.Embedding(
num_atoms + 1, embedding_dim, padding_idx=0
)
self.graph_token = nn.Embedding(1, embedding_dim)
self.degree_encoder = DegreeEncoder(
max_degree=max_degree, embedding_dim=embedding_dim
)
self.path_encoder = PathEncoder(
max_len=multi_hop_max_dist,
feat_dim=edge_dim,
num_heads=num_attention_heads,
)
self.spatial_encoder = SpatialEncoder(
max_dist=num_spatial, num_heads=num_attention_heads
)
self.graph_token_virtual_distance = nn.Embedding(1, num_attention_heads)
self.emb_layer_norm = nn.LayerNorm(self.embedding_dim)
self.layers = nn.ModuleList([])
self.layers.extend(
[
GraphormerLayer(
feat_size=self.embedding_dim,
hidden_size=ffn_embedding_dim,
num_heads=num_attention_heads,
dropout=dropout,
activation=activation_fn,
norm_first=pre_layernorm,
)
for _ in range(num_encoder_layers)
]
)
# map graph_rep to num_classes
self.lm_head_transform_weight = nn.Linear(
self.embedding_dim, self.embedding_dim
)
self.layer_norm = nn.LayerNorm(self.embedding_dim)
self.activation_fn = activation_fn
self.embed_out = nn.Linear(self.embedding_dim, num_classes, bias=False)
self.lm_output_learned_bias = nn.Parameter(th.zeros(num_classes))
def reset_output_layer_parameters(self):
self.lm_output_learned_bias = nn.Parameter(th.zeros(1))
self.embed_out.reset_parameters()
def forward(
self,
node_feat,
in_degree,
out_degree,
path_data,
dist,
attn_mask=None,
):
num_graphs, max_num_nodes, _ = node_feat.shape
deg_emb = self.degree_encoder(th.stack((in_degree, out_degree)))
# node feature + degree encoding as input
node_feat = self.atom_encoder(node_feat.int()).sum(dim=-2)
node_feat = node_feat + deg_emb
graph_token_feat = self.graph_token.weight.unsqueeze(0).repeat(
num_graphs, 1, 1
)
x = th.cat([graph_token_feat, node_feat], dim=1)
# spatial encoding and path encoding serve as attention bias
attn_bias = th.zeros(
num_graphs,
max_num_nodes + 1,
max_num_nodes + 1,
self.num_heads,
device=dist.device,
)
path_encoding = self.path_encoder(dist, path_data)
spatial_encoding = self.spatial_encoder(dist)
attn_bias[:, 1:, 1:, :] = path_encoding + spatial_encoding
# spatial encoding of the virtual node
t = self.graph_token_virtual_distance.weight.reshape(
1, 1, self.num_heads
)
# Since the virtual node comes first, the spatial encodings between it
# and other nodes will fill the 1st row and 1st column (omit num_graphs
# and num_heads dimensions) of attn_bias matrix by broadcasting.
attn_bias[:, 1:, 0, :] = attn_bias[:, 1:, 0, :] + t
attn_bias[:, 0, :, :] = attn_bias[:, 0, :, :] + t
x = self.emb_layer_norm(x)
for layer in self.layers:
x = layer(
x,
attn_mask=attn_mask,
attn_bias=attn_bias,
)
graph_rep = x[:, 0, :]
graph_rep = self.layer_norm(
self.activation_fn(self.lm_head_transform_weight(graph_rep))
)
graph_rep = self.embed_out(graph_rep) + self.lm_output_learned_bias
return graph_rep