wehub/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
examples/core/gat/README.md
+24
View File
@@ -0,0 +1,24 @@
Graph Attention Networks (GAT)
============
- Paper link: [https://arxiv.org/abs/1710.10903](https://arxiv.org/abs/1710.10903)
- Author's code repo (tensorflow implementation):
[https://github.com/PetarV-/GAT](https://github.com/PetarV-/GAT).
- Popular pytorch implementation:
[https://github.com/Diego999/pyGAT](https://github.com/Diego999/pyGAT).
How to run
-------
Run with the following for multiclass node classification (available datasets: "cora", "citeseer", "pubmed")
```bash
python3 train.py --dataset cora
```
> **_NOTE:_** Users may occasionally run into low accuracy issue (e.g., test accuracy < 0.8) due to overfitting. This can be resolved by adding Early Stopping or reducing maximum number of training epochs.
Summary
-------
* cora: ~0.821
* citeseer: ~0.710
* pubmed: ~0.780
examples/core/gat/train.py
+140
View File
@@ -0,0 +1,140 @@
import argparse
import time
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl import AddSelfLoop
from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
class GAT(nn.Module):
def __init__(self, in_size, hid_size, out_size, heads):
super().__init__()
self.gat_layers = nn.ModuleList()
# two-layer GAT
self.gat_layers.append(
dglnn.GATConv(
in_size,
hid_size,
heads[0],
feat_drop=0.6,
attn_drop=0.6,
activation=F.elu,
)
)
self.gat_layers.append(
dglnn.GATConv(
hid_size * heads[0],
out_size,
heads[1],
feat_drop=0.6,
attn_drop=0.6,
activation=None,
)
)
def forward(self, g, inputs):
h = inputs
for i, layer in enumerate(self.gat_layers):
h = layer(g, h)
if i == len(self.gat_layers) - 1: # last layer
h = h.mean(1)
else: # other layer(s)
h = h.flatten(1)
return h
def evaluate(g, features, labels, mask, model):
model.eval()
with torch.no_grad():
logits = model(g, features)
logits = logits[mask]
labels = labels[mask]
_, indices = torch.max(logits, dim=1)
correct = torch.sum(indices == labels)
return correct.item() * 1.0 / len(labels)
def train(g, features, labels, masks, model, num_epochs):
# Define train/val samples, loss function and optimizer
train_mask = masks[0]
val_mask = masks[1]
loss_fcn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=5e-3, weight_decay=5e-4)
for epoch in range(num_epochs):
t0 = time.time()
model.train()
logits = model(g, features)
loss = loss_fcn(logits[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
acc = evaluate(g, features, labels, val_mask, model)
t1 = time.time()
print(
"Epoch {:05d} | Loss {:.4f} | Accuracy {:.4f} | Time {:.4f}".format(
epoch, loss.item(), acc, t1 - t0
)
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--dataset",
type=str,
default="cora",
help="Dataset name ('cora', 'citeseer', 'pubmed').",
)
parser.add_argument(
"--num_epochs",
type=int,
default=200,
help="Number of epochs for train.",
)
parser.add_argument(
"--num_gpus",
type=int,
default=0,
help="Number of GPUs used for train and evaluation.",
)
args = parser.parse_args()
print(f"Training with DGL built-in GATConv module.")
# Load and preprocess dataset
transform = (
AddSelfLoop()
) # by default, it will first remove self-loops to prevent duplication
if args.dataset == "cora":
data = CoraGraphDataset(transform=transform)
elif args.dataset == "citeseer":
data = CiteseerGraphDataset(transform=transform)
elif args.dataset == "pubmed":
data = PubmedGraphDataset(transform=transform)
else:
raise ValueError("Unknown dataset: {}".format(args.dataset))
g = data[0]
if args.num_gpus > 0 and torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
g = g.int().to(device)
features = g.ndata["feat"]
labels = g.ndata["label"]
masks = g.ndata["train_mask"], g.ndata["val_mask"], g.ndata["test_mask"]
# Create GAT model
in_size = features.shape[1]
out_size = data.num_classes
model = GAT(in_size, 8, out_size, heads=[8, 1]).to(device)
print("Training...")
train(g, features, labels, masks, model, args.num_epochs)
print("Testing...")
acc = evaluate(g, features, labels, masks[2], model)
print("Test accuracy {:.4f}".format(acc))
examples/core/gated_gcn/README.md
+23
View File
@@ -0,0 +1,23 @@
Gated Graph ConvNet (GatedGCN)
==============================
* paper link: [https://arxiv.org/abs/2003.00982.pdf](https://arxiv.org/abs/2003.00982.pdf)
## Dataset
Task: Graph Property Prediction
| Dataset | #Graphs | #Node Feats | #Edge Feats | Metric |
| :---------: | :-----: | :---------: | :---------: | :-----: |
| ogbg-molhiv | 41,127 | 9 | 3 | ROC-AUC |
How to run
----------
```bash
python3 train.py --dataset ogbg-molhiv --num_gpus 0 --num_epochs 50
```
## Summary
* ogbg-molhiv: ~0.781
examples/core/gated_gcn/train.py
+167
View File
@@ -0,0 +1,167 @@
"""
Gated Graph Convolutional Network module for graph classification tasks
"""
import argparse
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from dgl.dataloading import GraphDataLoader
from dgl.nn.pytorch import GatedGCNConv
from dgl.nn.pytorch.glob import AvgPooling
from ogb.graphproppred import DglGraphPropPredDataset, Evaluator
from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
class GatedGCN(nn.Module):
def __init__(
self,
hid_dim,
out_dim,
num_layers,
dropout=0.2,
batch_norm=True,
residual=True,
activation=F.relu,
):
super(GatedGCN, self).__init__()
self.num_layers = num_layers
self.dropout = dropout
self.node_encoder = AtomEncoder(hid_dim)
self.edge_encoder = BondEncoder(hid_dim)
self.layers = nn.ModuleList()
for _ in range(self.num_layers):
layer = GatedGCNConv(
input_feats=hid_dim,
edge_feats=hid_dim,
output_feats=hid_dim,
dropout=dropout,
batch_norm=batch_norm,
residual=residual,
activation=activation,
)
self.layers.append(layer)
self.pooling = AvgPooling()
self.output = nn.Linear(hid_dim, out_dim)
def forward(self, g, node_feat, edge_feat):
# Encode node and edge feature.
hv = self.node_encoder(node_feat)
he = self.edge_encoder(edge_feat)
# GatedGCNConv layers.
for layer in self.layers:
hv, he = layer(g, hv, he)
# Output project.
h_g = self.pooling(g, hv)
return self.output(h_g)
def train(model, device, data_loader, opt, loss_fn):
model.train()
train_loss = []
for g, labels in data_loader:
g = g.to(device)
labels = labels.to(torch.float32).to(device)
logits = model(g, g.ndata["feat"], g.edata["feat"])
loss = loss_fn(logits, labels)
opt.zero_grad()
loss.backward()
opt.step()
train_loss.append(loss.item())
return sum(train_loss) / len(train_loss)
@torch.no_grad()
def evaluate(model, device, data_loader, evaluator):
model.eval()
y_true, y_pred = [], []
for g, labels in data_loader:
g = g.to(device)
logits = model(g, g.ndata["feat"], g.edata["feat"])
y_true.append(labels.detach().cpu())
y_pred.append(logits.detach().cpu())
y_true = torch.cat(y_true, dim=0).numpy()
y_pred = torch.cat(y_pred, dim=0).numpy()
return evaluator.eval({"y_true": y_true, "y_pred": y_pred})["rocauc"]
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--dataset",
type=str,
default="ogbg-molhiv",
help="Dataset name ('ogbg-molhiv', 'ogbg-molbace', 'ogbg-molmuv').",
)
parser.add_argument(
"--num_epochs",
type=int,
default=200,
help="Number of epochs for train.",
)
parser.add_argument(
"--num_gpus",
type=int,
default=0,
help="Number of GPUs used for train and evaluation.",
)
args = parser.parse_args()
print("Training with DGL built-in GATConv module.")
# Load ogb dataset & evaluator.
dataset = DglGraphPropPredDataset(name=args.dataset)
evaluator = Evaluator(name=args.dataset)
if args.num_gpus > 0 and torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
n_classes = dataset.num_tasks
split_idx = dataset.get_idx_split()
train_loader = GraphDataLoader(
dataset[split_idx["train"]],
batch_size=32,
shuffle=True,
)
valid_loader = GraphDataLoader(dataset[split_idx["valid"]], batch_size=32)
test_loader = GraphDataLoader(dataset[split_idx["test"]], batch_size=32)
# Load model.
model = GatedGCN(hid_dim=256, out_dim=n_classes, num_layers=8).to(device)
print(model)
opt = optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.BCEWithLogitsLoss()
print("---------- Training ----------")
for epoch in range(args.num_epochs):
# Kick off training.
t0 = time.time()
loss = train(model, device, train_loader, opt, loss_fn)
t1 = time.time()
# Evaluate the prediction.
val_acc = evaluate(model, device, valid_loader, evaluator)
print(
f"Epoch {epoch:05d} | Loss {loss:.4f} | Accuracy {val_acc:.4f} | "
f"Time {t1 - t0:.4f}"
)
acc = evaluate(model, device, test_loader, evaluator)
print(f"Test accuracy {acc:.4f}")
examples/core/graphsage/node_classification.py
+140
View File
@@ -0,0 +1,140 @@
"""
This script trains and tests a GraphSAGE model based on the information of
a full graph.
This flowchart describes the main functional sequence of the provided example.
main
│
├───> Load and preprocess full dataset
│
├───> Instantiate SAGE model
│
├───> train
│ │
│ └───> Training loop
│ │
│ └───> SAGE.forward
└───> test
│
└───> Evaluate the model
"""
import argparse
import time
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl import AddSelfLoop
from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
class SAGE(nn.Module):
def __init__(self, in_size, hidden_size, out_size):
super().__init__()
self.layers = nn.ModuleList()
# Two-layer GraphSAGE-gcn.
self.layers.append(dglnn.SAGEConv(in_size, hidden_size, "gcn"))
self.layers.append(dglnn.SAGEConv(hidden_size, out_size, "gcn"))
self.dropout = nn.Dropout(0.5)
def forward(self, graph, x):
hidden_x = x
for layer_idx, layer in enumerate(self.layers):
hidden_x = layer(graph, hidden_x)
is_last_layer = layer_idx == len(self.layers) - 1
if not is_last_layer:
hidden_x = F.relu(hidden_x)
hidden_x = self.dropout(hidden_x)
return hidden_x
def evaluate(g, features, labels, mask, model):
model.eval()
with torch.no_grad():
logits = model(g, features)
logits = logits[mask]
labels = labels[mask]
_, indices = torch.max(logits, dim=1)
correct = torch.sum(indices == labels)
return correct.item() * 1.0 / len(labels)
def train(g, features, labels, masks, model):
# Define train/val samples, loss function and optimizer.
train_mask, val_mask = masks
loss_fcn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
# Training loop.
for epoch in range(200):
t0 = time.time()
model.train()
logits = model(g, features)
loss = loss_fcn(logits[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
t1 = time.time()
acc = evaluate(g, features, labels, val_mask, model)
print(
f"Epoch {epoch:05d} | Loss {loss.item():.4f} | Accuracy {acc:.4f} | "
f"Time {t1 - t0:.4f}"
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="GraphSAGE")
parser.add_argument(
"--dataset",
type=str,
default="cora",
help="Dataset name ('cora', 'citeseer', 'pubmed')",
)
args = parser.parse_args()
print(f"Training with DGL built-in GraphSage module")
#####################################################################
# (HIGHLIGHT) Node classification task is a supervise learning task
# in which the model try to predict the label of a certain node.
# In this example, graph sage algorithm is applied to this task.
# A good accuracy can be achieved after a few steps of training.
#
# First, the whole graph is loaded and transformed. Then the training
# process is performed on a model which is composed of 2 GraphSAGE-gcn
# layer. Finally, the performance of the model is evaluated on test set.
#####################################################################
# Load and preprocess dataset.
transform = (
AddSelfLoop()
) # By default, it will first remove self-loops to prevent duplication.
if args.dataset == "cora":
data = CoraGraphDataset(transform=transform)
elif args.dataset == "citeseer":
data = CiteseerGraphDataset(transform=transform)
elif args.dataset == "pubmed":
data = PubmedGraphDataset(transform=transform)
else:
raise ValueError(f"Unknown dataset: {args.dataset}")
g = data[0]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
g = g.int().to(device)
features = g.ndata["feat"]
labels = g.ndata["label"]
masks = (g.ndata["train_mask"], g.ndata["val_mask"])
# Create GraphSAGE model.
in_size = features.shape[1]
out_size = data.num_classes
model = SAGE(in_size, 16, out_size).to(device)
# Model training.
print("Training...")
train(g, features, labels, masks, model)
# Test the model.
print("Testing...")
acc = evaluate(g, features, labels, g.ndata["test_mask"], model)
print(f"Test accuracy {acc:.4f}")
examples/core/rgcn/README.md
+65
View File
@@ -0,0 +1,65 @@
# Node classification on heterogeneous graph with RGCN
This example aims to demonstrate how to run node classification task on heterogeneous graph with **DGL**. Models are not tuned to achieve the best accuracy yet.
## Run on `ogbn-mag` dataset
In the preprocess stage, reverse edges are added and duplicate edges are removed. Feature data of `author` and `institution` node types are generated dynamically with embedding layer.
### Sample on CPU and train/infer on CPU
```
python3 hetero_rgcn.py --dataset ogbn-mag
```
### Sample on CPU and train/infer on GPU
```
python3 hetero_rgcn.py --dataset ogbn-mag --num_gpus 1
```
### Resource usage and time cost
Below results are roughly collected from an AWS EC2 **g4dn.metal**, 384GB RAM, 96 vCPUs(Cascade Lake P-8259L), 8 NVIDIA T4 GPUs(16GB RAM). CPU RAM usage is the peak value of `used` field of `free` command which is a bit rough. Please refer to `RSS`/`USS`/`PSS` which are more accurate. GPU RAM usage is the peak value recorded by `nvidia-smi` command.
| Dataset Size | CPU RAM Usage | Num of GPUs | GPU RAM Usage | Time Per Epoch(Training) |
| ------------ | ------------- | ----------- | ------------- | ------------------------ |
| ~1.1GB | ~7GB | 0 | 0GB | ~233s |
| ~1.1GB | ~5GB | 1 | 4.5GB | ~73.6s |
### Accuracies
```
Epoch: 01, Loss: 2.3386, Valid: 47.67%, Test: 46.96%
Epoch: 02, Loss: 1.5563, Valid: 47.66%, Test: 47.02%
Epoch: 03, Loss: 1.1557, Valid: 46.58%, Test: 45.42%
Test accuracy 45.3850
```
## Run on `ogb-lsc-mag240m` dataset
In the preprocess stage, reverse edges are added and duplicate edges are removed. What's more, feature data are generated in advance for `author` and `institution` node types via message passing. Since such preprocessing will usually take a long time, we also offer the above files for download:
* [`paper-feat.npy`](https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/paper-feat.npy)
* [`author-feat.npy`](https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/author-feat.npy)
* [`inst-feat.npy`](https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/inst-feat.npy)
* [`hetero-graph.dgl`](https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/hetero-graph.dgl)
### Sample on CPU and train/infer on CPU
```
python3 hetero_rgcn.py --dataset ogb-lsc-mag240m
```
### Sample on CPU and train/infer on GPU
```
python3 hetero_rgcn.py --dataset ogb-lsc-mag240m --num_gpus 1
```
### Resource usage and time cost
Below results are roughly collected from an AWS EC2 **g4dn.metal**, 384GB RAM, 96 vCPUs(Cascade Lake P-8259L), 8 NVIDIA T4 GPUs(16GB RAM). CPU RAM usage is the peak value of `used` field of `free` command which is a bit rough. Please refer to `RSS`/`USS`/`PSS` which are more accurate. GPU RAM usage is the peak value recorded by `nvidia-smi` command.
| Dataset Size | CPU RAM Usage | Num of GPUs | GPU RAM Usage | Time Per Epoch(Training) |
| ------------ | ------------- | ----------- | ------------- | ------------------------ |
| ~404GB | ~72GB | 0 | 0GB | ~325s |
| ~404GB | ~61GB | 1 | 14GB | ~178s |
### Accuracies
```
Epoch: 01, Loss: 2.0798, Valid: 52.04%
Epoch: 02, Loss: 1.8652, Valid: 54.51%
Epoch: 03, Loss: 1.8175, Valid: 53.71%
```
examples/core/rgcn/hetero_rgcn.py
+686
View File
@@ -0,0 +1,686 @@
"""
This script, `hetero_rgcn.py`, trains and tests a Relational Graph
Convolutional Network (R-GCN) model for node classification on the
Open Graph Benchmark (OGB) dataset "ogbn-mag". For more details on
"ogbn-mag", please refer to the OGB website:
(https://ogb.stanford.edu/docs/linkprop/)
Paper [Modeling Relational Data with Graph Convolutional Networks]
(https://arxiv.org/abs/1703.06103).
Generation of graph embeddings is the main difference between homograph
node classification and heterograph node classification:
- Homograph: Since all nodes and edges are of the same type, embeddings
can be generated using a unified approach. Type-specific handling is
typically not required.
- Heterograph: Due to the existence of multiple types of nodes and edges,
specific embeddings need to be generated for each type. This allows for
a more nuanced capture of the complex structure and semantic information
within the heterograph.
This flowchart describes the main functional sequence of the provided example.
main
│
├───> prepare_data
│ │
│ └───> Load and preprocess dataset
│
├───> rel_graph_embed [HIGHLIGHT]
│ │
│ └───> Generate graph embeddings
│
├───> Instantiate RGCN model
│ │
│ ├───> RelGraphConvLayer (input to hidden)
│ │
│ └───> RelGraphConvLayer (hidden to output)
│
└───> train
│
│
└───> Training loop
│
├───> EntityClassify.forward (RGCN model forward pass)
│
└───> test
│
└───> EntityClassify.evaluate
"""
import argparse
import itertools
import sys
import time
import dgl
import dgl.nn as dglnn
import numpy as np
import psutil
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl import AddReverse, Compose, ToSimple
from dgl.nn import HeteroEmbedding
from ogb.lsc import MAG240MDataset, MAG240MEvaluator
from ogb.nodeproppred import DglNodePropPredDataset, Evaluator
from tqdm import tqdm
def prepare_data(args, device):
feats = {}
if args.dataset == "ogbn-mag":
dataset = DglNodePropPredDataset(name="ogbn-mag", root=args.rootdir)
# - graph: dgl graph object.
# - label: torch tensor of shape (num_nodes, num_tasks).
g, labels = dataset[0]
# Flatten the labels for "paper" type nodes. This step reduces the
# dimensionality of the labels. We need to flatten the labels because
# the model requires a 1-dimensional label tensor.
labels = labels["paper"].flatten().long()
# Apply transformation to the graph.
# - "ToSimple()" removes multi-edge between two nodes.
# - "AddReverse()" adds reverse edges to the graph.
print("Start to transform graph. This may take a while...")
transform = Compose([ToSimple(), AddReverse()])
g = transform(g)
else:
dataset = MAG240MDataset(root=args.rootdir)
(g,), _ = dgl.load_graphs(args.graph_path)
g = g.formats(["csc"])
labels = torch.as_tensor(dataset.paper_label).long()
# As feature data is too large to fit in memory, we read it from disk.
feats["paper"] = torch.as_tensor(
np.load(args.paper_feature_path, mmap_mode="r+")
)
feats["author"] = torch.as_tensor(
np.load(args.author_feature_path, mmap_mode="r+")
)
feats["institution"] = torch.as_tensor(
np.load(args.inst_feature_path, mmap_mode="r+")
)
print(f"Loaded graph: {g}")
# Get train/valid/test index.
split_idx = dataset.get_idx_split()
if args.dataset == "ogb-lsc-mag240m":
split_idx = {
split_type: {"paper": split_idx[split_type]}
for split_type in split_idx
}
# Initialize a train sampler that samples neighbors for multi-layer graph
# convolution. It samples 25 and 10 neighbors for the first and second
# layers respectively.
sampler = dgl.dataloading.MultiLayerNeighborSampler([25, 10], fused=False)
num_workers = args.num_workers
train_loader = dgl.dataloading.DataLoader(
g,
split_idx["train"],
sampler,
batch_size=1024,
shuffle=True,
num_workers=num_workers,
device=device,
)
return g, labels, dataset.num_classes, split_idx, train_loader, feats
def extract_embed(node_embed, input_nodes):
emb = node_embed(
{ntype: input_nodes[ntype] for ntype in input_nodes if ntype != "paper"}
)
return emb
def rel_graph_embed(graph, embed_size):
"""Initialize a heterogenous embedding layer for all node types in the
graph, except for the "paper" node type.
The function constructs a dictionary 'node_num', where the keys are node
types (ntype) and the values are the number of nodes for each type. This
dictionary is used to create a HeteroEmbedding instance.
(HIGHLIGHT)
A HeteroEmbedding instance holds separate embedding layers for each node
type, each with its own feature space of dimensionality
(node_num[ntype], embed_size), where 'node_num[ntype]' is the number of
nodes of type 'ntype' and 'embed_size' is the embedding dimension.
The "paper" node type is specifically excluded, possibly because these nodes
might already have predefined feature representations, and therefore, do not
require an additional embedding layer.
Parameters
----------
graph : DGLGraph
The graph for which to create the heterogenous embedding layer.
embed_size : int
The size of the embedding vectors.
Returns
--------
HeteroEmbedding
A heterogenous embedding layer for all node types in the graph, except
for the "paper" node type.
"""
node_num = {}
for ntype in graph.ntypes:
# Skip the "paper" node type.
if ntype == "paper":
continue
node_num[ntype] = graph.num_nodes(ntype)
return HeteroEmbedding(node_num, embed_size)
class RelGraphConvLayer(nn.Module):
def __init__(
self,
in_size,
out_size,
ntypes,
relation_names,
activation=None,
dropout=0.0,
):
super(RelGraphConvLayer, self).__init__()
self.in_size = in_size
self.out_size = out_size
self.ntypes = ntypes
self.relation_names = relation_names
self.activation = activation
########################################################################
# (HIGHLIGHT) HeteroGraphConv is a graph convolution operator over
# heterogeneous graphs. A dictionary is passed where the key is the
# relation name and the value is the instance of GraphConv. norm="right"
# is to divide the aggregated messages by each node’s in-degrees, which
# is equivalent to averaging the received messages. weight=False and
# bias=False as we will use our own weight matrices defined later.
########################################################################
self.conv = dglnn.HeteroGraphConv(
{
rel: dglnn.GraphConv(
in_size, out_size, norm="right", weight=False, bias=False
)
for rel in relation_names
}
)
# Create a separate Linear layer for each relationship. Each
# relationship has its own weights which will be applied to the node
# features before performing convolution.
self.weight = nn.ModuleDict(
{
rel_name: nn.Linear(in_size, out_size, bias=False)
for rel_name in self.relation_names
}
)
# Create a separate Linear layer for each node type.
# loop_weights are used to update the output embedding of each target node
# based on its own features, thereby allowing the model to refine the node
# representations. Note that this does not imply the existence of self-loop
# edges in the graph. It is similar to residual connection.
self.loop_weights = nn.ModuleDict(
{
ntype: nn.Linear(in_size, out_size, bias=True)
for ntype in self.ntypes
}
)
self.loop_weights = nn.ModuleDict(
{
ntype: nn.Linear(in_size, out_size, bias=True)
for ntype in self.ntypes
}
)
self.dropout = nn.Dropout(dropout)
# Initialize parameters of the model.
self.reset_parameters()
def reset_parameters(self):
for layer in self.weight.values():
layer.reset_parameters()
for layer in self.loop_weights.values():
layer.reset_parameters()
def forward(self, g, inputs):
"""
Parameters
----------
g : DGLGraph
Input graph.
inputs : dict[str, torch.Tensor]
Node feature for each node type.
Returns
-------
dict[str, torch.Tensor]
New node features for each node type.
"""
# Create a deep copy of the graph g with features saved in local
# frames to prevent side effects from modifying the graph.
g = g.local_var()
# Create a dictionary of weights for each relationship. The weights
# are retrieved from the Linear layers defined earlier.
weight_dict = {
rel_name: {"weight": self.weight[rel_name].weight.T}
for rel_name in self.relation_names
}
# Create a dictionary of node features for the destination nodes in
# the graph. We slice the node features according to the number of
# destination nodes of each type. This is necessary because when
# incorporating the effect of self-loop edges, we perform computations
# only on the destination nodes' features. By doing so, we ensure the
# feature dimensions match and prevent any misuse of incorrect node
# features.
inputs_dst = {
k: v[: g.number_of_dst_nodes(k)] for k, v in inputs.items()
}
# Apply the convolution operation on the graph. mod_kwargs are
# additional arguments for each relation function defined in the
# HeteroGraphConv. In this case, it's the weights for each relation.
hs = self.conv(g, inputs, mod_kwargs=weight_dict)
def _apply(ntype, h):
# Apply the `loop_weight` to the input node features, effectively
# acting as a residual connection. This allows the model to refine
# node embeddings based on its current features.
h = h + self.loop_weights[ntype](inputs_dst[ntype])
if self.activation:
h = self.activation(h)
return self.dropout(h)
# Apply the function defined above for each node type. This will update
# the node features using the `loop_weights`, apply the activation
# function and dropout.
return {ntype: _apply(ntype, h) for ntype, h in hs.items()}
class EntityClassify(nn.Module):
def __init__(self, g, in_size, out_size):
super(EntityClassify, self).__init__()
self.in_size = in_size
self.hidden_size = 64
self.out_size = out_size
# Generate and sort a list of unique edge types from the input graph.
# eg. ['writes', 'cites']
self.relation_names = list(set(g.etypes))
self.relation_names.sort()
self.dropout = 0.5
self.layers = nn.ModuleList()
# First layer: transform input features to hidden features. Use ReLU
# as the activation function and apply dropout for regularization.
self.layers.append(
RelGraphConvLayer(
self.in_size,
self.hidden_size,
g.ntypes,
self.relation_names,
activation=F.relu,
dropout=self.dropout,
)
)
# Second layer: transform hidden features to output features. No
# activation function is applied at this stage.
self.layers.append(
RelGraphConvLayer(
self.hidden_size,
self.out_size,
g.ntypes,
self.relation_names,
activation=None,
)
)
def reset_parameters(self):
# Reset the parameters of each layer.
for layer in self.layers:
layer.reset_parameters()
def forward(self, h, blocks):
for layer, block in zip(self.layers, blocks):
h = layer(block, h)
return h
def extract_node_features(name, g, input_nodes, node_embed, feats, device):
"""Extract the node features from embedding layer or raw features."""
if name == "ogbn-mag":
# Extract node embeddings for the input nodes.
node_features = extract_embed(node_embed, input_nodes)
# Add the batch's raw "paper" features. Corresponds to the content
# in the function `rel_graph_embed` comment.
node_features.update(
{"paper": g.ndata["feat"]["paper"][input_nodes["paper"].cpu()]}
)
node_features = {k: e.to(device) for k, e in node_features.items()}
else:
node_features = {
ntype: feats[ntype][input_nodes[ntype].cpu()].to(device)
for ntype in input_nodes
}
# Original feature data are stored in float16 while model weights are
# float32, so we need to convert the features to float32.
# [TODO] Enable mixed precision training on GPU.
node_features = {k: v.float() for k, v in node_features.items()}
return node_features
def train(
dataset,
g,
feats,
model,
node_embed,
optimizer,
train_loader,
split_idx,
labels,
device,
):
print("Start training...")
category = "paper"
# Typically, the best Validation performance is obtained after
# the 1st or 2nd epoch. This is why the max epoch is set to 3.
for epoch in range(3):
num_train = split_idx["train"][category].shape[0]
t0 = time.time()
model.train()
total_loss = 0
for input_nodes, seeds, blocks in tqdm(
train_loader, desc=f"Epoch {epoch:02d}"
):
# Move the input data onto the device.
blocks = [blk.to(device) for blk in blocks]
# We only predict the nodes with type "category".
seeds = seeds[category]
batch_size = seeds.shape[0]
# Extract the node features from embedding layer or raw features.
node_features = extract_node_features(
dataset, g, input_nodes, node_embed, feats, device
)
lbl = labels[seeds.cpu()].to(device)
# Reset gradients.
optimizer.zero_grad()
# Generate predictions.
logits = model(node_features, blocks)[category]
y_hat = logits.log_softmax(dim=-1)
loss = F.nll_loss(y_hat, lbl)
loss.backward()
optimizer.step()
total_loss += loss.item() * batch_size
t1 = time.time()
loss = total_loss / num_train
# Evaluate the model on the val/test set.
valid_acc = evaluate(
dataset,
g,
feats,
model,
node_embed,
labels,
device,
split_idx["valid"],
)
test_key = "test" if dataset == "ogbn-mag" else "test-dev"
test_acc = evaluate(
dataset,
g,
feats,
model,
node_embed,
labels,
device,
split_idx[test_key],
save_test_submission=(dataset == "ogb-lsc-mag240m"),
)
print(
f"Epoch: {epoch +1 :02d}, "
f"Loss: {loss:.4f}, "
f"Valid: {100 * valid_acc:.2f}%, "
f"Test: {100 * test_acc:.2f}%, "
f"Time {t1 - t0:.4f}"
)
@torch.no_grad()
def evaluate(
dataset,
g,
feats,
model,
node_embed,
labels,
device,
idx,
save_test_submission=False,
):
# Switches the model to evaluation mode.
model.eval()
category = "paper"
if dataset == "ogbn-mag":
evaluator = Evaluator(name="ogbn-mag")
else:
evaluator = MAG240MEvaluator()
sampler = dgl.dataloading.MultiLayerNeighborSampler([25, 10], fused=False)
dataloader = dgl.dataloading.DataLoader(
g,
idx,
sampler,
batch_size=4096,
shuffle=False,
num_workers=0,
device=device,
)
# To store the predictions.
y_hats = list()
y_true = list()
for input_nodes, seeds, blocks in tqdm(dataloader, desc="Inference"):
blocks = [blk.to(device) for blk in blocks]
# We only predict the nodes with type "category".
node_features = extract_node_features(
dataset, g, input_nodes, node_embed, feats, device
)
# Generate predictions.
logits = model(node_features, blocks)[category]
# Apply softmax to the logits and get the prediction by selecting the
# argmax.
y_hat = logits.log_softmax(dim=-1).argmax(dim=1, keepdims=True)
y_hats.append(y_hat.cpu())
y_true.append(labels[seeds["paper"].cpu()])
y_pred = torch.cat(y_hats, dim=0)
y_true = torch.cat(y_true, dim=0)
y_true = torch.unsqueeze(y_true, 1)
if dataset == "ogb-lsc-mag240m":
y_pred = y_pred.view(-1)
y_true = y_true.view(-1)
if save_test_submission:
evaluator.save_test_submission(
input_dict={"y_pred": y_pred}, dir_path=".", mode="test-dev"
)
return evaluator.eval({"y_true": y_true, "y_pred": y_pred})["acc"]
def main(args):
device = (
"cuda:0" if torch.cuda.is_available() and args.num_gpus > 0 else "cpu"
)
# Prepare the data.
g, labels, num_classes, split_idx, train_loader, feats = prepare_data(
args, device
)
feat_size = 128 if args.dataset == "ogbn-mag" else 768
# Create the embedding layer and move it to the appropriate device.
embed_layer = None
if args.dataset == "ogbn-mag":
embed_layer = rel_graph_embed(g, feat_size).to(device)
print(
"Number of embedding parameters: "
f"{sum(p.numel() for p in embed_layer.parameters())}"
)
# Initialize the entity classification model.
model = EntityClassify(g, feat_size, num_classes).to(device)
print(
"Number of model parameters: "
f"{sum(p.numel() for p in model.parameters())}"
)
try:
if embed_layer is not None:
embed_layer.reset_parameters()
model.reset_parameters()
except:
# Old pytorch version doesn't support reset_parameters() API.
##################################################################
# [Why we need to reset the parameters?]
# If parameters are not reset, the model will start with the
# parameters learned from the last run, potentially resulting
# in biased outcomes or sub-optimal performance if the model was
# previously stuck in a poor local minimum.
##################################################################
pass
# `itertools.chain()` is a function in Python's itertools module.
# It is used to flatten a list of iterables, making them act as
# one big iterable.
# In this context, the following code is used to create a single
# iterable over the parameters of both the model and the embed_layer,
# which is passed to the optimizer. The optimizer then updates all
# these parameters during the training process.
all_params = itertools.chain(
model.parameters(),
[] if embed_layer is None else embed_layer.parameters(),
)
optimizer = torch.optim.Adam(all_params, lr=0.01)
# `expected_max`` is the number of physical cores on your machine.
# The `logical` parameter, when set to False, ensures that the count
# returned is the number of physical cores instead of logical cores
# (which could be higher due to technologies like Hyper-Threading).
expected_max = int(psutil.cpu_count(logical=False))
if args.num_workers >= expected_max:
print(
"[ERROR] You specified num_workers are larger than physical"
f"cores, please set any number less than {expected_max}",
file=sys.stderr,
)
train(
args.dataset,
g,
feats,
model,
embed_layer,
optimizer,
train_loader,
split_idx,
labels,
device,
)
print("Testing...")
test_key = "test" if args.dataset == "ogbn-mag" else "test-dev"
test_acc = evaluate(
args.dataset,
g,
feats,
model,
embed_layer,
labels,
device,
split_idx[test_key],
save_test_submission=(args.dataset == "ogb-lsc-mag240m"),
)
print(f"Test accuracy {test_acc*100:.4f}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="RGCN")
parser.add_argument(
"--dataset",
type=str,
default="ogbn-mag",
help="Dataset for train: ogbn-mag, ogb-lsc-mag240m",
)
parser.add_argument(
"--num_gpus",
type=int,
default=0,
help="Number of GPUs. Use 0 for CPU training.",
)
parser.add_argument(
"--num_workers",
type=int,
default=0,
help="Number of worker processes for data loading.",
)
parser.add_argument(
"--rootdir",
type=str,
default="./dataset/",
help="Directory to download the OGB dataset.",
)
parser.add_argument(
"--graph_path",
type=str,
default="./graph.dgl",
help="Path to the graph file.",
)
parser.add_argument(
"--paper_feature_path",
type=str,
default="./paper-feat.npy",
help="Path to the features of paper nodes.",
)
parser.add_argument(
"--author_feature_path",
type=str,
default="./author-feat.npy",
help="Path to the features of author nodes.",
)
parser.add_argument(
"--inst_feature_path",
type=str,
default="./inst-feat.npy",
help="Path to the features of institution nodes.",
)
args = parser.parse_args()
main(args)