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/graphbolt/README.md
+11
View File
@@ -0,0 +1,11 @@
## How to run the code?
```bash
python link_prediction.py
```
Results (10 epochs):
```
Valid MRR 0.7040
Test MRR 0.7043
```
examples/graphbolt/disk_based_feature/README.md
+110
View File
@@ -0,0 +1,110 @@
## Overview
This project demonstrates how to use GraphBolt to train and evaluate a GraphSAGE model for node classification task on large graphs, where node features are on-disk and fetched using `DiskBasedFeature`. GraphBolt utilizes various in-house implemented caching policy algorithms such as [SIEVE](https://cachemon.github.io/SIEVE-website/), [S3-FIFO](https://s3fifo.com), LRU and [CLOCK](https://people.csail.mit.edu/saltzer/Multics/MHP-Saltzer-060508/bookcases/M00s/M0104%20074-12%29.PDF) to cache frequently required features and io_uring to fetch cache-missed features from disk. The SIEVE algorithm is the default option.
# Node classification task
This example demonstrates how to run node classification task with **GraphBolt.DiskBasedFeature**. All results are collected on an AWS EC2 g5.8xlarge instance with 128GB RAM, 32 cores, an 24GB A10G GPU and a instance storage of 250K IOPS.
## Run on `ogbn-papers100M` dataset
| Dataset | Graph Size | Feature Size | Feature Dim |
| :-------------: | :--------: | :----------: | :---------: |
| ogbn-papers100M | 13 GB | 53 GB | 128 |
## Results with various caching policies
This part trains a three-layer GraphSAGE model for 3 epochs on `ogbn-papers100M` dataset with 10GB CPU cache, using neighbor sampling.
### Run default SIEVE policy
Instruction:
```
python node_classification.py --gpu-cache-size-in-gigabytes=0 --cpu-cache-size-in-gigabytes=10 --dataset=ogbn-papers100M --epochs=3
```
Result:
```
Training: 1178it [03:00, 6.53it/s, num_nodes=671260, gpu_cache_miss=1, cpu_cache_miss=0.0578]
Evaluating: 123it [00:16, 7.47it/s, num_nodes=624816, gpu_cache_miss=1, cpu_cache_miss=0.0569]
Epoch 00, Loss: 1.4173, Approx. Train: 0.5787, Approx. Val: 0.6353, Time: 180.33928060531616s
Training: 1178it [01:39, 11.79it/s, num_nodes=648380, gpu_cache_miss=1, cpu_cache_miss=0.0451]
Evaluating: 123it [00:15, 7.90it/s, num_nodes=625373, gpu_cache_miss=1, cpu_cache_miss=0.0451]
Epoch 01, Loss: 1.1446, Approx. Train: 0.6386, Approx. Val: 0.6382, Time: 99.92613315582275s
Training: 1178it [01:36, 12.15it/s, num_nodes=674194, gpu_cache_miss=1, cpu_cache_miss=0.0408]
Evaluating: 123it [00:15, 8.08it/s, num_nodes=628233, gpu_cache_miss=1, cpu_cache_miss=0.0409]
Epoch 02, Loss: 1.0975, Approx. Train: 0.6507, Approx. Val: 0.6535, Time: 96.95083212852478s
```
### Performance Comparison on four caching polices
Below results demonstrate the epoch time with four different caching policies.
| Policy | Epoch 1 (s) | Epoch 2 (s) | Epoch 3 (s) |
| :-----: | :---------: | :---------: | :---------: |
| SIEVE | 180.339 | 99.926 | 96.951 |
| S3-FiFO | 181.438 | 110.054 | 108.310 |
| LRU | 194.583 | 138.352 | 138.369 |
| CLOCK | 188.915 | 129.372 | 129.388 |
## Results with Layer-Neighbor Sampling
This part trains a three-layer GraphSAGE model for 3 epochs on `ogbn-papers100M` dataset with 10GB CPU cache, using Layer-Neighbor Sampling and default SIEVE policy.
### Run default `--batch-dependency=1`
Instruction:
```
python node_classification.py --gpu-cache-size-in-gigabytes=0 --cpu-cache-size-in-gigabytes=10 --dataset=ogbn-papers100M --sample-mode=sample_layer_neighbor --batch-dependency=1 --epochs=3
```
Result:
```
Training: 1178it [02:51, 6.88it/s, num_nodes=463495, gpu_cache_miss=1, cpu_cache_miss=0.0774]
Evaluating: 123it [00:15, 7.94it/s, num_nodes=465592, gpu_cache_miss=1, cpu_cache_miss=0.0762]
Epoch 00, Loss: 1.4173, Approx. Train: 0.5774, Approx. Val: 0.6300, Time: 171.11454963684082s
Training: 1178it [01:34, 12.43it/s, num_nodes=474446, gpu_cache_miss=1, cpu_cache_miss=0.0604]
Evaluating: 123it [00:14, 8.45it/s, num_nodes=462042, gpu_cache_miss=1, cpu_cache_miss=0.0603]
Epoch 01, Loss: 1.1463, Approx. Train: 0.6384, Approx. Val: 0.6395, Time: 94.7821741104126s
Training: 1178it [01:31, 12.82it/s, num_nodes=479331, gpu_cache_miss=1, cpu_cache_miss=0.0545]
Evaluating: 123it [00:14, 8.67it/s, num_nodes=463628, gpu_cache_miss=1, cpu_cache_miss=0.0546]
Epoch 02, Loss: 1.1000, Approx. Train: 0.6501, Approx. Val: 0.6516, Time: 91.8746063709259s
```
### Performance Comparison on different `--batch-dependency`
| batch-dependency | Epoch 1 (s) | Epoch 2 (s) | Epoch 3 (s) |
| :--------------: | :---------: | :---------: | :---------: |
| 1 | 171.114 | 94.782 | 91.875 |
| 64 | 144.241 | 78.749 | 75.270 |
| 4096 | 92.494 | 56.111 | 57.647 |
### Effect of `--layer-dependency`
Below results demonstrate the effect of enabling `--layer-dependency` on epoch time when setting `--batch-dependency=1`.
| layer-dependency | Epoch 1 (s) | Epoch 2 (s) | Epoch 3 (s) |
| :--------------: | :---------: | :---------: | :---------: |
| False | 171.114 | 94.782 | 91.875 |
| True | 159.625 | 86.209 | 83.171 |
## Compared to In-mem Performance
This part trains a three-layer GraphSAGE model for 3 epochs on `ogbn-papers100M` dataset with 20GB CPU cache and 5GB GPU cache, using neighbor sampling. We compare it to the in-mem performance with 5GB GPU cache. Following result demonstrates that with sufficient cache memory, the performance of DiskBasedFeature is not bottlenecked by the cache itself and comparable with in-memory feature stores. Note that the first epoch of training initiates the cache, thus taking longer time.
Instruction:
```
python node_classification.py --gpu-cache-size-in-gigabytes=5 --cpu-cache-size-in-gigabytes=20 --dataset=ogbn-papers100M --epochs=3
```
Result:
| Feature Store | Epoch 1 (s) | Epoch 2 (s) | Epoch 3 (s) |
| :--------------: | :---------: | :---------: | :---------: |
| DiskBasedFeature | 143.761 | 32.018 | 31.889 |
| In-memory | 28.861 | 28.330 | 28.305 |
examples/graphbolt/disk_based_feature/node_classification.py
+540
View File
@@ -0,0 +1,540 @@
"""
This example references examples/graphbolt/pyg/labor/node_classification.py
"""
import argparse
import time
from copy import deepcopy
import dgl.graphbolt as gb
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
def accuracy(out, labels):
assert out.ndim == 2
assert out.size(0) == labels.size(0)
assert labels.ndim == 1 or (labels.ndim == 2 and labels.size(1) == 1)
labels = labels.flatten()
predictions = torch.argmax(out, 1)
return (labels == predictions).sum(dtype=torch.float64) / labels.size(0)
class SAGE(nn.Module):
def __init__(self, in_size, hidden_size, out_size, num_layers, dropout):
super().__init__()
self.layers = nn.ModuleList()
# Three-layer GraphSAGE-mean.
self.layers.append(dglnn.SAGEConv(in_size, hidden_size, "mean"))
for _ in range(num_layers - 2):
self.layers.append(dglnn.SAGEConv(hidden_size, hidden_size, "mean"))
self.layers.append(dglnn.SAGEConv(hidden_size, out_size, "mean"))
self.dropout = nn.Dropout(dropout)
self.hidden_size = hidden_size
self.out_size = out_size
# Set the dtype for the layers manually.
self.set_layer_dtype(torch.float32)
def set_layer_dtype(self, _dtype):
for layer in self.layers:
for param in layer.parameters():
param.data = param.data.to(_dtype)
def forward(self, blocks, x):
hidden_x = x
for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)):
hidden_x = layer(block, 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 inference(self, graph, features, dataloader, storage_device):
"""Conduct layer-wise inference to get all the node embeddings."""
pin_memory = storage_device == "pinned"
buffer_device = torch.device("cpu" if pin_memory else storage_device)
for layer_idx, layer in enumerate(self.layers):
is_last_layer = layer_idx == len(self.layers) - 1
y = torch.empty(
graph.total_num_nodes,
self.out_size if is_last_layer else self.hidden_size,
dtype=torch.float32,
device=buffer_device,
pin_memory=pin_memory,
)
for data in tqdm(dataloader):
# len(blocks) = 1
hidden_x = layer(data.blocks[0], data.node_features["feat"])
if not is_last_layer:
hidden_x = F.relu(hidden_x)
hidden_x = self.dropout(hidden_x)
# By design, our output nodes are contiguous.
y[data.seeds[0] : data.seeds[-1] + 1] = hidden_x.to(
buffer_device
)
if not is_last_layer:
features.update("node", None, "feat", y)
return y
def create_dataloader(
graph, features, itemset, batch_size, fanout, device, job
):
# Initialize an ItemSampler to sample mini-batches from the dataset.
datapipe = gb.ItemSampler(
itemset,
batch_size=batch_size,
shuffle=(job == "train"),
drop_last=(job == "train"),
)
# Copy the data to the specified device.
if args.graph_device != "cpu":
datapipe = datapipe.copy_to(device=device)
# Sample neighbors for each node in the mini-batch.
kwargs = (
{
# Layer dependency makes it so that the sampled neighborhoods across layers
# become correlated, reducing the total number of sampled unique nodes in a
# minibatch, thus reducing the amount of feature data requested.
"layer_dependency": args.layer_dependency,
# Batch dependency makes it so that the sampled neighborhoods across minibatches
# become correlated, reducing the total number of sampled unique nodes across
# minibatches, thus increasing temporal locality and reducing cache miss rates.
"batch_dependency": args.batch_dependency,
}
if args.sample_mode == "sample_layer_neighbor"
else {}
)
datapipe = getattr(datapipe, args.sample_mode)(
graph,
fanout if job != "infer" else [-1],
overlap_fetch=args.overlap_graph_fetch,
**kwargs,
)
# Copy the data to the specified device.
if args.feature_device != "cpu":
datapipe = datapipe.copy_to(device=device)
# Fetch node features for the sampled subgraph.
datapipe = datapipe.fetch_feature(
features,
node_feature_keys=["feat"],
overlap_fetch=args.overlap_feature_fetch,
)
# Copy the data to the specified device.
if args.feature_device == "cpu":
datapipe = datapipe.copy_to(device=device)
# Create and return a DataLoader to handle data loading.
return gb.DataLoader(datapipe, num_workers=args.num_workers)
def train_step(minibatch, optimizer, model, loss_fn):
node_features = minibatch.node_features["feat"]
labels = minibatch.labels
optimizer.zero_grad()
out = model(minibatch.blocks, node_features)
loss = loss_fn(out, labels)
num_correct = accuracy(out, labels) * labels.size(0)
loss.backward()
optimizer.step()
return loss.detach(), num_correct, labels.size(0)
def train_helper(
dataloader,
model,
optimizer,
loss_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
):
model.train() # Set the model to training mode
total_loss = torch.zeros(1, device=device) # Accumulator for the total loss
# Accumulator for the total number of correct predictions
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0 # Accumulator for the total number of samples processed
num_batches = 0 # Counter for the number of mini-batches processed
start = time.time()
dataloader = tqdm(dataloader, "Training")
for step, minibatch in enumerate(dataloader):
loss, num_correct, num_samples = train_step(
minibatch, optimizer, model, loss_fn
)
total_loss += loss
total_correct += num_correct
total_samples += num_samples
num_batches += 1
if step % 25 == 0:
# log every 25 steps for performance.
dataloader.set_postfix(
{
"num_nodes": minibatch.node_ids().size(0),
"gpu_cache_miss": gpu_cache_miss_rate_fn(),
"cpu_cache_miss": cpu_cache_miss_rate_fn(),
}
)
train_loss = total_loss / num_batches
train_acc = total_correct / total_samples
end = time.time()
return train_loss, train_acc, end - start
def train(
train_dataloader,
valid_dataloader,
model,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
):
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
loss_fn = nn.CrossEntropyLoss()
best_model = None
best_model_acc = 0
best_model_epoch = -1
for epoch in range(args.epochs):
train_loss, train_acc, duration = train_helper(
train_dataloader,
model,
optimizer,
loss_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
)
val_acc = evaluate(
model,
valid_dataloader,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
)
if val_acc > best_model_acc:
best_model_acc = val_acc
best_model = deepcopy(model.state_dict())
best_model_epoch = epoch
print(
f"Epoch {epoch:02d}, Loss: {train_loss.item():.4f}, "
f"Approx. Train: {train_acc.item():.4f}, "
f"Approx. Val: {val_acc.item():.4f}, "
f"Time: {duration}s"
)
if best_model_epoch + args.early_stopping_patience < epoch:
break
return best_model
@torch.no_grad()
def layerwise_infer(
args,
graph,
features,
itemsets,
all_nodes_set,
model,
):
model.eval()
dataloader = create_dataloader(
graph=graph,
features=features,
itemset=all_nodes_set,
batch_size=args.batch_size,
fanout=[-1],
device=args.device,
job="infer",
)
pred = model.inference(graph, features, dataloader, args.feature_device)
metrics = {}
for split_name, itemset in itemsets.items():
nid, labels = itemset[:]
acc = accuracy(
pred[nid.to(pred.device)],
labels.to(pred.device),
)
metrics[split_name] = acc.item()
return metrics
def evaluate_step(minibatch, model):
node_features = minibatch.node_features["feat"]
labels = minibatch.labels
out = model(minibatch.blocks, node_features)
num_correct = accuracy(out, labels) * labels.size(0)
return num_correct, labels.size(0)
@torch.no_grad()
def evaluate(
model,
dataloader,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
):
model.eval()
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0
val_dataloader_tqdm = tqdm(dataloader, "Evaluating")
for step, minibatch in enumerate(val_dataloader_tqdm):
num_correct, num_samples = evaluate_step(minibatch, model)
total_correct += num_correct
total_samples += num_samples
if step % 25 == 0:
val_dataloader_tqdm.set_postfix(
{
"num_nodes": minibatch.node_ids().size(0),
"gpu_cache_miss": gpu_cache_miss_rate_fn(),
"cpu_cache_miss": cpu_cache_miss_rate_fn(),
}
)
return total_correct / total_samples
def parse_args():
parser = argparse.ArgumentParser(
description="Which dataset are you going to use?"
)
parser.add_argument(
"--epochs", type=int, default=9999999, help="Number of training epochs."
)
parser.add_argument(
"--lr",
type=float,
default=0.001,
help="Learning rate for optimization.",
)
parser.add_argument("--num-hidden", type=int, default=256)
parser.add_argument("--dropout", type=float, default=0.2)
parser.add_argument(
"--batch-size", type=int, default=1024, help="Batch size for training."
)
parser.add_argument(
"--num-workers",
type=int,
default=0,
help="Number of workers for data loading.",
)
parser.add_argument(
"--dataset",
type=str,
default="ogbn-products",
choices=[
"ogbn-arxiv",
"ogbn-products",
"ogbn-papers100M",
"igb-hom-tiny",
"igb-hom-small",
"igb-hom-medium",
"igb-hom-large",
"igb-hom",
],
)
parser.add_argument("--root", type=str, default="datasets")
parser.add_argument(
"--fanout",
type=str,
default="10,10,10",
help="Fan-out of neighbor sampling. len(fanout) determines the number of"
" GNN layers in your model. Default: 10,10,10",
)
parser.add_argument(
"--mode",
default="pinned-pinned-cuda",
choices=[
"cpu-cpu-cpu",
"cpu-cpu-cuda",
"cpu-pinned-cuda",
"pinned-pinned-cuda",
"cuda-pinned-cuda",
"cuda-cuda-cuda",
],
help="Graph storage - feature storage - Train device: 'cpu' for CPU and"
" RAM, 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
)
parser.add_argument("--layer-dependency", action="store_true")
parser.add_argument("--batch-dependency", type=int, default=1)
parser.add_argument(
"--cpu-feature-cache-policy",
type=str,
default=None,
choices=["s3-fifo", "sieve", "lru", "clock"],
help="The cache policy for the CPU feature cache.",
)
parser.add_argument(
"--cpu-cache-size-in-gigabytes",
type=float,
default=0,
help="The capacity of the CPU cache in GiB.",
)
parser.add_argument(
"--gpu-cache-size-in-gigabytes",
type=float,
default=0,
help="The capacity of the GPU cache in GiB.",
)
parser.add_argument("--early-stopping-patience", type=int, default=25)
parser.add_argument(
"--sample-mode",
default="sample_neighbor",
choices=["sample_neighbor", "sample_layer_neighbor"],
help="The sampling function when doing layerwise sampling.",
)
parser.add_argument("--precision", type=str, default="high")
parser.add_argument("--enable-inference", action="store_true")
return parser.parse_args()
def main():
torch.set_float32_matmul_precision(args.precision)
if not torch.cuda.is_available():
args.mode = "cpu-cpu-cpu"
print(f"Training in {args.mode} mode.")
args.graph_device, args.feature_device, args.device = args.mode.split("-")
args.overlap_feature_fetch = args.feature_device == "pinned"
args.overlap_graph_fetch = args.graph_device == "pinned"
"""
Load and preprocess on-disk dataset.
We inspect the in_memory field of the feature_data in the YAML file and modify
it to False. This will make sure the feature_data is loaded as DiskBasedFeature.
"""
print("Loading data...")
disk_based_feature_keys = None
if args.cpu_cache_size_in_gigabytes > 0:
disk_based_feature_keys = [("node", None, "feat")]
dataset = gb.BuiltinDataset(args.dataset, root=args.root)
if disk_based_feature_keys is None:
disk_based_feature_keys = set()
for feature in dataset.yaml_data["feature_data"]:
feature_key = (feature["domain"], feature["type"], feature["name"])
# Set the in_memory setting to False without modifying YAML file.
if feature_key in disk_based_feature_keys:
feature["in_memory"] = False
dataset = dataset.load()
# Move the dataset to the selected storage.
graph = (
dataset.graph.pin_memory_()
if args.graph_device == "pinned"
else dataset.graph.to(args.graph_device)
)
features = (
dataset.feature.pin_memory_()
if args.feature_device == "pinned"
else dataset.feature.to(args.feature_device)
)
train_set = dataset.tasks[0].train_set
valid_set = dataset.tasks[0].validation_set
test_set = dataset.tasks[0].test_set
all_nodes_set = dataset.all_nodes_set
args.fanout = list(map(int, args.fanout.split(",")))
num_classes = dataset.tasks[0].metadata["num_classes"]
"""
If the CPU cache size is greater than 0, we wrap the DiskBasedFeature to be
a CPUCachedFeature. This internally manages the CPU feature cache by the
specified cache replacement policy. This will reduce the amount of data
transferred during disk read operations for this feature.
Note: It is advised to set the CPU cache size to be at least 4 times the number
of sampled nodes in a mini-batch, otherwise the feature fetcher might get into
a deadlock, causing a hang.
"""
if args.cpu_cache_size_in_gigabytes > 0 and isinstance(
features[("node", None, "feat")], gb.DiskBasedFeature
):
features[("node", None, "feat")] = gb.cpu_cached_feature(
features[("node", None, "feat")],
int(args.cpu_cache_size_in_gigabytes * 1024 * 1024 * 1024),
args.cpu_feature_cache_policy,
args.feature_device == "pinned",
)
cpu_cached_feature = features[("node", None, "feat")]
cpu_cache_miss_rate_fn = lambda: cpu_cached_feature.miss_rate
else:
cpu_cache_miss_rate_fn = lambda: 1
"""
If the GPU cache size is greater than 0, we wrap the underlying feature store
to be a GPUCachedFeature. This will reduce the amount of data transferred during
host-to-device copy operations for this feature.
"""
if args.gpu_cache_size_in_gigabytes > 0 and args.feature_device != "cuda":
features[("node", None, "feat")] = gb.gpu_cached_feature(
features[("node", None, "feat")],
int(args.gpu_cache_size_in_gigabytes * 1024 * 1024 * 1024),
)
gpu_cached_feature = features[("node", None, "feat")]
gpu_cache_miss_rate_fn = lambda: gpu_cached_feature.miss_rate
else:
gpu_cache_miss_rate_fn = lambda: 1
train_dataloader, valid_dataloader = (
create_dataloader(
graph=graph,
features=features,
itemset=itemset,
batch_size=args.batch_size,
fanout=args.fanout,
device=args.device,
job=job,
)
for itemset, job in zip([train_set, valid_set], ["train", "evaluate"])
)
in_channels = features.size("node", None, "feat")[0]
model = SAGE(
in_channels,
args.num_hidden,
num_classes,
len(args.fanout),
args.dropout,
).to(args.device)
assert len(args.fanout) == len(model.layers)
best_model = train(
train_dataloader,
valid_dataloader,
model,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
args.device,
)
model.load_state_dict(best_model)
if args.enable_inference:
# Test the model.
print("Testing...")
itemsets = {"train": train_set, "val": valid_set, "test": test_set}
final_acc = layerwise_infer(
args,
graph,
features,
itemsets,
all_nodes_set,
model,
)
print("Final accuracy values:")
print(final_acc)
if __name__ == "__main__":
args = parse_args()
print(args)
main()
examples/graphbolt/lightning/README.md
+13
View File
@@ -0,0 +1,13 @@
# Node classification on homogeneous graph with GraphSAGE
## Run on `ogbn-products` dataset
### Command
```
python3 node_classification.py
```
### Results
```
Valid Accuracy: 0.907
```
examples/graphbolt/lightning/node_classification.py
+233
View File
@@ -0,0 +1,233 @@
"""
This flowchart describes the main functional sequence of the provided example.
main
│
├───> Instantiate DataModule
│ │
│ └───> Load dataset
│ │
│ └───> Create train and valid dataloader[HIGHLIGHT]
│ │
│ └───> ItemSampler (Distribute data to minibatchs)
│ │
│ └───> sample_neighbor or sample_layer_neighbor
(Sample a subgraph for a minibatch)
│ │
│ └───> fetch_feature (Fetch features for the sampled subgraph)
│
├───> Instantiate GraphSAGE model
│ │
│ ├───> SAGEConvLayer (input to hidden)
│ │
│ └───> SAGEConvLayer (hidden to hidden)
│ │
│ └───> SAGEConvLayer (hidden to output)
│ │
│ └───> DropoutLayer
│
└───> Run
│
│
└───> Trainer[HIGHLIGHT]
│
├───> SAGE.forward (GraphSAGE model forward pass)
│
└───> Validate
"""
import argparse
import dgl.graphbolt as gb
import dgl.nn.pytorch as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
from pytorch_lightning import LightningDataModule, LightningModule, Trainer
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from torchmetrics import Accuracy
class SAGE(LightningModule):
def __init__(self, in_feats, n_hidden, n_classes):
super().__init__()
self.save_hyperparameters()
self.layers = nn.ModuleList()
self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, "mean"))
self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, "mean"))
self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, "mean"))
self.dropout = nn.Dropout(0.5)
self.n_hidden = n_hidden
self.n_classes = n_classes
self.train_acc = Accuracy(task="multiclass", num_classes=n_classes)
self.val_acc = Accuracy(task="multiclass", num_classes=n_classes)
def forward(self, blocks, x):
h = x
for l, (layer, block) in enumerate(zip(self.layers, blocks)):
h = layer(block, h)
if l != len(self.layers) - 1:
h = F.relu(h)
h = self.dropout(h)
return h
def log_node_and_edge_counts(self, blocks):
node_counts = [block.num_src_nodes() for block in blocks] + [
blocks[-1].num_dst_nodes()
]
edge_counts = [block.num_edges() for block in blocks]
for i, c in enumerate(node_counts):
self.log(
f"num_nodes/{i}",
float(c),
prog_bar=True,
on_step=True,
on_epoch=False,
)
if i < len(edge_counts):
self.log(
f"num_edges/{i}",
float(edge_counts[i]),
prog_bar=True,
on_step=True,
on_epoch=False,
)
def training_step(self, batch, batch_idx):
blocks = [block.to("cuda") for block in batch.blocks]
x = batch.node_features["feat"]
y = batch.labels.to("cuda")
y_hat = self(blocks, x)
loss = F.cross_entropy(y_hat, y)
self.train_acc(torch.argmax(y_hat, 1), y)
self.log(
"train_acc",
self.train_acc,
prog_bar=True,
on_step=True,
on_epoch=False,
)
self.log_node_and_edge_counts(blocks)
return loss
def validation_step(self, batch, batch_idx):
blocks = [block.to("cuda") for block in batch.blocks]
x = batch.node_features["feat"]
y = batch.labels.to("cuda")
y_hat = self(blocks, x)
self.val_acc(torch.argmax(y_hat, 1), y)
self.log(
"val_acc",
self.val_acc,
prog_bar=True,
on_step=False,
on_epoch=True,
sync_dist=True,
)
self.log_node_and_edge_counts(blocks)
def configure_optimizers(self):
optimizer = torch.optim.Adam(
self.parameters(), lr=0.001, weight_decay=5e-4
)
return optimizer
class DataModule(LightningDataModule):
def __init__(self, dataset, fanouts, batch_size, num_workers):
super().__init__()
self.fanouts = fanouts
self.batch_size = batch_size
self.num_workers = num_workers
self.feature_store = dataset.feature
self.graph = dataset.graph
self.train_set = dataset.tasks[0].train_set
self.valid_set = dataset.tasks[0].validation_set
self.num_classes = dataset.tasks[0].metadata["num_classes"]
def create_dataloader(self, node_set, is_train):
datapipe = gb.ItemSampler(
node_set,
batch_size=self.batch_size,
shuffle=True,
drop_last=True,
)
sampler = (
datapipe.sample_layer_neighbor
if is_train
else datapipe.sample_neighbor
)
datapipe = sampler(self.graph, self.fanouts)
datapipe = datapipe.fetch_feature(self.feature_store, ["feat"])
dataloader = gb.DataLoader(datapipe, num_workers=self.num_workers)
return dataloader
########################################################################
# (HIGHLIGHT) The 'train_dataloader' and 'val_dataloader' hooks are
# essential components of the Lightning framework, defining how data is
# loaded during training and validation. In this example, we utilize a
# specialized 'graphbolt dataloader', which are concatenated by a series
# of datapipes, for these purposes.
########################################################################
def train_dataloader(self):
return self.create_dataloader(self.train_set, is_train=True)
def val_dataloader(self):
return self.create_dataloader(self.valid_set, is_train=False)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="GNN baselines on ogbn-products data with GraphBolt"
)
parser.add_argument(
"--num_gpus",
type=int,
default=1,
help="number of GPUs used for computing (default: 1)",
)
parser.add_argument(
"--batch_size",
type=int,
default=1024,
help="input batch size for training (default: 1024)",
)
parser.add_argument(
"--epochs",
type=int,
default=40,
help="number of epochs to train (default: 40)",
)
parser.add_argument(
"--num_workers",
type=int,
default=0,
help="number of workers (default: 0)",
)
args = parser.parse_args()
dataset = gb.BuiltinDataset("ogbn-products").load()
datamodule = DataModule(
dataset,
[10, 10, 10],
args.batch_size,
args.num_workers,
)
in_size = dataset.feature.size("node", None, "feat")[0]
model = SAGE(in_size, 256, datamodule.num_classes)
# Train.
checkpoint_callback = ModelCheckpoint(monitor="val_acc", mode="max")
early_stopping_callback = EarlyStopping(monitor="val_acc", mode="max")
########################################################################
# (HIGHLIGHT) The `Trainer` is the key Class in lightning, which automates
# everything after defining `LightningDataModule` and
# `LightningDataModule`. More details can be found in
# https://lightning.ai/docs/pytorch/stable/common/trainer.html.
########################################################################
trainer = Trainer(
accelerator="gpu",
devices=args.num_gpus,
max_epochs=args.epochs,
callbacks=[checkpoint_callback, early_stopping_callback],
)
trainer.fit(model, datamodule=datamodule)
examples/graphbolt/link_prediction.py
+428
View File
@@ -0,0 +1,428 @@
"""
This script trains and tests a GraphSAGE model for link prediction on
large graphs using graphbolt dataloader.
Paper: [Inductive Representation Learning on Large Graphs]
(https://arxiv.org/abs/1706.02216)
Unlike previous dgl examples, we've utilized the newly defined dataloader
from GraphBolt. This example will help you grasp how to build an end-to-end
training pipeline using GraphBolt.
While node classification predicts labels for nodes based on their
local neighborhoods, link prediction assesses the likelihood of an edge
existing between two nodes, necessitating different sampling strategies
that account for pairs of nodes and their joint neighborhoods.
TODO: Add the link_prediction.py example to core/graphsage.
Before reading this example, please familiar yourself with graphsage link
prediction by reading the example in the
`examples/core/graphsage/link_prediction.py`
If you want to train graphsage on a large graph in a distributed fashion, read
the example in the `examples/distributed/graphsage/`.
This flowchart describes the main functional sequence of the provided example.
main
│
├───> OnDiskDataset pre-processing
│
├───> Instantiate SAGE model
│
├───> train
│ │
│ ├───> Get graphbolt dataloader (HIGHLIGHT)
│ │
│ └───> Training loop
│ │
│ ├───> SAGE.forward
│ │
│ └───> Validation set evaluation
│
└───> Test set evaluation
"""
import argparse
import time
from functools import partial
import dgl.graphbolt as gb
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
import tqdm
from torchmetrics.retrieval import RetrievalMRR
class SAGE(nn.Module):
def __init__(self, in_size, hidden_size):
super().__init__()
self.layers = nn.ModuleList()
self.layers.append(dglnn.SAGEConv(in_size, hidden_size, "mean"))
self.layers.append(dglnn.SAGEConv(hidden_size, hidden_size, "mean"))
self.layers.append(dglnn.SAGEConv(hidden_size, hidden_size, "mean"))
self.hidden_size = hidden_size
self.predictor = nn.Sequential(
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, 1),
)
def forward(self, blocks, x):
hidden_x = x
for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)):
hidden_x = layer(block, hidden_x)
is_last_layer = layer_idx == len(self.layers) - 1
if not is_last_layer:
hidden_x = F.relu(hidden_x)
return hidden_x
def inference(self, graph, features, dataloader, storage_device):
"""Conduct layer-wise inference to get all the node embeddings."""
pin_memory = storage_device == "pinned"
buffer_device = torch.device("cpu" if pin_memory else storage_device)
print("Start node embedding inference.")
for layer_idx, layer in enumerate(self.layers):
is_last_layer = layer_idx == len(self.layers) - 1
y = torch.empty(
graph.total_num_nodes,
self.hidden_size,
dtype=torch.float32,
device=buffer_device,
pin_memory=pin_memory,
)
for data in tqdm.tqdm(dataloader):
# len(blocks) = 1
hidden_x = layer(data.blocks[0], data.node_features["feat"])
if not is_last_layer:
hidden_x = F.relu(hidden_x)
# By design, our seed nodes are contiguous.
y[data.seeds[0] : data.seeds[-1] + 1] = hidden_x.to(
buffer_device, non_blocking=True
)
if not is_last_layer:
features.update("node", None, "feat", y)
return y
def create_dataloader(args, graph, features, itemset, is_train=True):
"""Get a GraphBolt version of a dataloader for link prediction tasks. This
function demonstrates how to utilize functional forms of datapipes in
GraphBolt. Alternatively, you can create a datapipe using its class
constructor. For a more detailed tutorial, please read the examples in
`dgl/notebooks/graphbolt/walkthrough.ipynb`.
"""
############################################################################
# [Input]:
# 'itemset': The current dataset.
# 'args.batch_size': Specify the number of samples to be processed together,
# referred to as a 'mini-batch'. (The term 'mini-batch' is used here to
# indicate a subset of the entire dataset that is processed together. This
# is in contrast to processing the entire dataset, known as a 'full batch'.)
# 'is_train': Determining if data should be shuffled. (Shuffling is
# generally used only in training to improve model generalization. It's
# not used in validation and testing as the focus there is to evaluate
# performance rather than to learn from the data.)
# [Output]:
# An ItemSampler object for handling mini-batch sampling.
# [Role]:
# Initialize the ItemSampler to sample mini-batche from the dataset.
############################################################################
datapipe = gb.ItemSampler(
itemset,
batch_size=args.train_batch_size if is_train else args.eval_batch_size,
shuffle=is_train,
)
############################################################################
# [Input]:
# 'device': The device to copy the data to.
# [Output]:
# A CopyTo object to copy the data to the specified device. Copying here
# ensures that the rest of the operations run on the GPU.
############################################################################
if args.storage_device != "cpu":
datapipe = datapipe.copy_to(device=args.device)
############################################################################
# [Input]:
# 'args.neg_ratio': Specify the ratio of negative to positive samples.
# (E.g., if neg_ratio is 1, for each positive sample there will be 1
# negative sample.)
# 'graph': The overall network topology for negative sampling.
# [Output]:
# A UniformNegativeSampler object that will handle the generation of
# negative samples for link prediction tasks.
# [Role]:
# Initialize the UniformNegativeSampler for negative sampling in link
# prediction.
# [Note]:
# If 'is_train' is False, the UniformNegativeSampler will not be used.
# Since, in validation and testing, the itemset already contains the
# negative edges information.
############################################################################
if is_train:
datapipe = datapipe.sample_uniform_negative(graph, args.neg_ratio)
############################################################################
# [Input]:
# 'datapipe' is either 'ItemSampler' or 'UniformNegativeSampler' depending
# on whether training is needed ('is_train'),
# 'graph': The network topology for sampling.
# 'args.fanout': Number of neighbors to sample per node.
# [Output]:
# A NeighborSampler object to sample neighbors.
# [Role]:
# Initialize a neighbor sampler for sampling the neighborhoods of nodes.
############################################################################
datapipe = datapipe.sample_neighbor(
graph,
args.fanout if is_train else [-1],
overlap_fetch=args.storage_device == "pinned",
asynchronous=args.storage_device != "cpu",
)
############################################################################
# [Input]:
# 'gb.exclude_seed_edges': Function to exclude seed edges, optionally
# including their reverse edges, from the sampled subgraphs in the
# minibatch.
# [Output]:
# A MiniBatchTransformer object with excluded seed edges.
# [Role]:
# During the training phase of link prediction, negative edges are
# sampled. It's essential to exclude the seed edges from the process
# to ensure that positive samples are not inadvertently included within
# the negative samples.
############################################################################
if is_train and args.exclude_edges:
datapipe = datapipe.exclude_seed_edges(
include_reverse_edges=True,
asynchronous=args.storage_device != "cpu",
)
############################################################################
# [Input]:
# 'features': The node features.
# 'node_feature_keys': The node feature keys (list) to be fetched.
# [Output]:
# A FeatureFetcher object to fetch node features.
# [Role]:
# Initialize a feature fetcher for fetching features of the sampled
# subgraphs.
############################################################################
datapipe = datapipe.fetch_feature(features, node_feature_keys=["feat"])
############################################################################
# [Input]:
# 'device': The device to copy the data to.
# [Output]:
# A CopyTo object to copy the data to the specified device.
############################################################################
if args.storage_device == "cpu":
datapipe = datapipe.copy_to(device=args.device)
############################################################################
# [Input]:
# 'datapipe': The datapipe object to be used for data loading.
# 'args.num_workers': The number of processes to be used for data loading.
# [Output]:
# A DataLoader object to handle data loading.
# [Role]:
# Initialize a multi-process dataloader to load the data in parallel.
############################################################################
dataloader = gb.DataLoader(
datapipe,
num_workers=args.num_workers,
)
# Return the fully-initialized DataLoader object.
return dataloader
@torch.no_grad()
def compute_mrr(args, model, node_emb, seeds, labels, indexes):
"""Compute the Mean Reciprocal Rank (MRR) for given source and destination
nodes.
This function computes the MRR for a set of node pairs, dividing the task
into batches to handle potentially large graphs.
"""
preds = torch.empty(seeds.shape[0], device=indexes.device)
mrr = RetrievalMRR()
seeds_src, seeds_dst = seeds.T
# The constant number is 1001, due to negtive ratio in the `ogbl-citation2`
# dataset is 1000.
eval_size = args.eval_batch_size * 1001
# Loop over node pairs in batches.
for start in tqdm.trange(0, seeds_src.shape[0], eval_size, desc="Evaluate"):
end = min(start + eval_size, seeds_src.shape[0])
# Fetch embeddings for current batch of source and destination nodes.
h_src = node_emb[seeds_src[start:end]].to(args.device)
h_dst = node_emb[seeds_dst[start:end]].to(args.device)
# Compute prediction scores using the model.
pred = model.predictor(h_src * h_dst).squeeze()
preds[start:end] = pred
return mrr(preds, labels, indexes=indexes)
@torch.no_grad()
def evaluate(args, model, graph, features, all_nodes_set, valid_set, test_set):
"""Evaluate the model on validation and test sets."""
model.eval()
dataloader = create_dataloader(
args, graph, features, all_nodes_set, is_train=False
)
# Compute node embeddings for the entire graph.
node_emb = model.inference(graph, features, dataloader, args.storage_device)
results = []
# Loop over both validation and test sets.
for split in [valid_set, test_set]:
# Unpack the item set.
seeds = split._items[0].to(node_emb.device)
labels = split._items[1].to(node_emb.device)
indexes = split._items[2].to(node_emb.device)
# Compute MRR values for the current split.
results.append(
compute_mrr(args, model, node_emb, seeds, labels, indexes)
)
return results
def train(args, model, graph, features, train_set):
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
dataloader = create_dataloader(args, graph, features, train_set)
for epoch in range(args.epochs):
model.train()
total_loss = 0
start_epoch_time = time.time()
for step, data in tqdm.tqdm(enumerate(dataloader)):
# Get node pairs with labels for loss calculation.
compacted_seeds = data.compacted_seeds.T
labels = data.labels
node_feature = data.node_features["feat"]
blocks = data.blocks
# Get the embeddings of the input nodes.
y = model(blocks, node_feature)
logits = model.predictor(
y[compacted_seeds[0]] * y[compacted_seeds[1]]
).squeeze()
# Compute loss.
loss = F.binary_cross_entropy_with_logits(logits, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
if step + 1 == args.early_stop:
break
end_epoch_time = time.time()
print(
f"Epoch {epoch:05d} | "
f"Loss {(total_loss) / (step + 1):.4f} | "
f"Time {(end_epoch_time - start_epoch_time):.4f} s"
)
def parse_args():
parser = argparse.ArgumentParser(description="OGBL-Citation2 (GraphBolt)")
parser.add_argument("--epochs", type=int, default=10)
parser.add_argument("--lr", type=float, default=0.0005)
parser.add_argument("--neg-ratio", type=int, default=1)
parser.add_argument("--train-batch-size", type=int, default=512)
parser.add_argument("--eval-batch-size", type=int, default=1024)
parser.add_argument("--num-workers", type=int, default=0)
parser.add_argument(
"--early-stop",
type=int,
default=0,
help="0 means no early stop, otherwise stop at the input-th step",
)
parser.add_argument(
"--fanout",
type=str,
default="15,10,5",
help="Fan-out of neighbor sampling. Default: 15,10,5",
)
parser.add_argument(
"--exclude-edges",
type=int,
default=1,
help="Whether to exclude reverse edges during sampling. Default: 1",
)
parser.add_argument(
"--mode",
default="pinned-cuda",
choices=["cpu-cpu", "cpu-cuda", "pinned-cuda", "cuda-cuda"],
help="Dataset storage placement and Train device: 'cpu' for CPU and RAM,"
" 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
)
return parser.parse_args()
def main(args):
if not torch.cuda.is_available():
args.mode = "cpu-cpu"
print(f"Training in {args.mode} mode.")
args.storage_device, args.device = args.mode.split("-")
args.device = torch.device(args.device)
# Load and preprocess dataset.
print("Loading data")
dataset = gb.BuiltinDataset("ogbl-citation2").load()
# Move the dataset to the selected storage.
if args.storage_device == "pinned":
graph = dataset.graph.pin_memory_()
features = dataset.feature.pin_memory_()
else:
graph = dataset.graph.to(args.storage_device)
features = dataset.feature.to(args.storage_device)
train_set = dataset.tasks[0].train_set
args.fanout = list(map(int, args.fanout.split(",")))
in_size = features.size("node", None, "feat")[0]
hidden_channels = 256
args.device = torch.device(args.device)
model = SAGE(in_size, hidden_channels).to(args.device)
# Model training.
print("Training...")
train(args, model, graph, features, train_set)
# Test the model.
print("Testing...")
test_set = dataset.tasks[0].test_set
valid_set = dataset.tasks[0].validation_set
all_nodes_set = dataset.all_nodes_set
valid_mrr, test_mrr = evaluate(
args, model, graph, features, all_nodes_set, valid_set, test_set
)
print(
f"Validation MRR {valid_mrr.item():.4f}, "
f"Test MRR {test_mrr.item():.4f}"
)
if __name__ == "__main__":
args = parse_args()
main(args)
examples/graphbolt/node_classification.py
+451
View File
@@ -0,0 +1,451 @@
"""
This script trains and tests a GraphSAGE model for node classification
on large graphs using GraphBolt dataloader.
Paper: [Inductive Representation Learning on Large Graphs]
(https://arxiv.org/abs/1706.02216)
Unlike previous dgl examples, we've utilized the newly defined dataloader
from GraphBolt. This example will help you grasp how to build an end-to-end
training pipeline using GraphBolt.
Before reading this example, please familiar yourself with graphsage node
classification by reading the example in the
`examples/core/graphsage/node_classification.py`. This introduction,
[A Blitz Introduction to Node Classification with DGL]
(https://docs.dgl.ai/tutorials/blitz/1_introduction.html), might be helpful.
If you want to train graphsage on a large graph in a distributed fashion,
please read the example in the `examples/distributed/graphsage/`.
This flowchart describes the main functional sequence of the provided example:
main
│
├───> OnDiskDataset pre-processing
│
├───> Instantiate SAGE model
│
├───> train
│ │
│ ├───> Get graphbolt dataloader (HIGHLIGHT)
│ │
│ └───> Training loop
│ │
│ ├───> SAGE.forward
│ │
│ └───> Validation set evaluation
│
└───> All nodes set inference & Test set evaluation
"""
import argparse
import time
import dgl.graphbolt as gb
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchmetrics.functional as MF
from tqdm import tqdm
def create_dataloader(
graph, features, itemset, batch_size, fanout, device, num_workers, job
):
"""
[HIGHLIGHT]
Get a GraphBolt version of a dataloader for node classification tasks.
This function demonstrates how to utilize functional forms of datapipes in
GraphBolt. For a more detailed tutorial, please read the examples in
`dgl/notebooks/graphbolt/walkthrough.ipynb`.
Alternatively, you can create a datapipe using its class constructor.
Parameters
----------
job : one of ["train", "evaluate", "infer"]
The stage where dataloader is created, with options "train", "evaluate"
and "infer".
Other parameters are explicated in the comments below.
"""
############################################################################
# [Step-1]:
# gb.ItemSampler()
# [Input]:
# 'itemset': The current dataset. (e.g. `train_set` or `valid_set`)
# 'batch_size': Specify the number of samples to be processed together,
# referred to as a 'mini-batch'. (The term 'mini-batch' is used here to
# indicate a subset of the entire dataset that is processed together. This
# is in contrast to processing the entire dataset, known as a 'full batch'.)
# 'job': Determines whether data should be shuffled. (Shuffling is
# generally used only in training to improve model generalization. It's
# not used in validation and testing as the focus there is to evaluate
# performance rather than to learn from the data.)
# [Output]:
# An ItemSampler object for handling mini-batch sampling.
# [Role]:
# Initialize the ItemSampler to sample mini-batche from the dataset.
############################################################################
datapipe = gb.ItemSampler(
itemset, batch_size=batch_size, shuffle=(job == "train")
)
############################################################################
# [Step-2]:
# self.copy_to()
# [Input]:
# 'device': The device to copy the data to.
# [Output]:
# A CopyTo object to copy the data to the specified device. Copying here
# ensures that the rest of the operations run on the GPU.
############################################################################
if args.storage_device != "cpu":
datapipe = datapipe.copy_to(device=device)
############################################################################
# [Step-3]:
# self.sample_neighbor()
# [Input]:
# 'graph': The network topology for sampling.
# '[-1] or fanout': Number of neighbors to sample per node. In
# training or validation, the length of `fanout` should be equal to the
# number of layers in the model. In inference, this parameter is set to
# [-1], indicating that all neighbors of a node are sampled.
# [Output]:
# A NeighborSampler object to sample neighbors.
# [Role]:
# Initialize a neighbor sampler for sampling the neighborhoods of nodes.
############################################################################
datapipe = getattr(datapipe, args.sample_mode)(
graph,
fanout if job != "infer" else [-1],
overlap_fetch=args.storage_device == "pinned",
asynchronous=args.storage_device != "cpu",
)
############################################################################
# [Step-4]:
# self.fetch_feature()
# [Input]:
# 'features': The node features.
# 'node_feature_keys': The keys of the node features to be fetched.
# [Output]:
# A FeatureFetcher object to fetch node features.
# [Role]:
# Initialize a feature fetcher for fetching features of the sampled
# subgraphs.
############################################################################
datapipe = datapipe.fetch_feature(features, node_feature_keys=["feat"])
############################################################################
# [Step-5]:
# self.copy_to()
# [Input]:
# 'device': The device to copy the data to.
# [Output]:
# A CopyTo object to copy the data to the specified device.
############################################################################
if args.storage_device == "cpu":
datapipe = datapipe.copy_to(device=device)
############################################################################
# [Step-6]:
# gb.DataLoader()
# [Input]:
# 'datapipe': The datapipe object to be used for data loading.
# 'num_workers': The number of processes to be used for data loading.
# [Output]:
# A DataLoader object to handle data loading.
# [Role]:
# Initialize a multi-process dataloader to load the data in parallel.
############################################################################
dataloader = gb.DataLoader(datapipe, num_workers=num_workers)
# Return the fully-initialized DataLoader object.
return dataloader
class SAGE(nn.Module):
def __init__(self, in_size, hidden_size, out_size):
super().__init__()
self.layers = nn.ModuleList()
# Three-layer GraphSAGE-mean.
self.layers.append(dglnn.SAGEConv(in_size, hidden_size, "mean"))
self.layers.append(dglnn.SAGEConv(hidden_size, hidden_size, "mean"))
self.layers.append(dglnn.SAGEConv(hidden_size, out_size, "mean"))
self.dropout = nn.Dropout(0.5)
self.hidden_size = hidden_size
self.out_size = out_size
# Set the dtype for the layers manually.
self.set_layer_dtype(torch.float32)
def set_layer_dtype(self, _dtype):
for layer in self.layers:
for param in layer.parameters():
param.data = param.data.to(_dtype)
def forward(self, blocks, x):
hidden_x = x
for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)):
hidden_x = layer(block, 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 inference(self, graph, features, dataloader, storage_device):
"""Conduct layer-wise inference to get all the node embeddings."""
pin_memory = storage_device == "pinned"
buffer_device = torch.device("cpu" if pin_memory else storage_device)
for layer_idx, layer in enumerate(self.layers):
is_last_layer = layer_idx == len(self.layers) - 1
y = torch.empty(
graph.total_num_nodes,
self.out_size if is_last_layer else self.hidden_size,
dtype=torch.float32,
device=buffer_device,
pin_memory=pin_memory,
)
for data in tqdm(dataloader):
# len(blocks) = 1
hidden_x = layer(data.blocks[0], data.node_features["feat"])
if not is_last_layer:
hidden_x = F.relu(hidden_x)
hidden_x = self.dropout(hidden_x)
# By design, our output nodes are contiguous.
y[data.seeds[0] : data.seeds[-1] + 1] = hidden_x.to(
buffer_device
)
if not is_last_layer:
features.update("node", None, "feat", y)
return y
@torch.no_grad()
def layerwise_infer(
args, graph, features, test_set, all_nodes_set, model, num_classes
):
model.eval()
dataloader = create_dataloader(
graph=graph,
features=features,
itemset=all_nodes_set,
batch_size=4 * args.batch_size,
fanout=[-1],
device=args.device,
num_workers=args.num_workers,
job="infer",
)
pred = model.inference(graph, features, dataloader, args.storage_device)
pred = pred[test_set._items[0]]
label = test_set._items[1].to(pred.device)
return MF.accuracy(
pred,
label,
task="multiclass",
num_classes=num_classes,
)
@torch.no_grad()
def evaluate(args, model, graph, features, itemset, num_classes):
model.eval()
y = []
y_hats = []
dataloader = create_dataloader(
graph=graph,
features=features,
itemset=itemset,
batch_size=args.batch_size,
fanout=args.fanout,
device=args.device,
num_workers=args.num_workers,
job="evaluate",
)
for step, data in tqdm(enumerate(dataloader), "Evaluating"):
x = data.node_features["feat"]
y.append(data.labels)
y_hats.append(model(data.blocks, x))
return MF.accuracy(
torch.cat(y_hats),
torch.cat(y),
task="multiclass",
num_classes=num_classes,
)
def train(args, graph, features, train_set, valid_set, num_classes, model):
optimizer = torch.optim.Adam(
model.parameters(), lr=args.lr, weight_decay=5e-4
)
dataloader = create_dataloader(
graph=graph,
features=features,
itemset=train_set,
batch_size=args.batch_size,
fanout=args.fanout,
device=args.device,
num_workers=args.num_workers,
job="train",
)
for epoch in range(args.epochs):
t0 = time.time()
model.train()
total_loss = 0
for step, data in tqdm(enumerate(dataloader), "Training"):
# The input features from the source nodes in the first layer's
# computation graph.
x = data.node_features["feat"]
# The ground truth labels from the destination nodes
# in the last layer's computation graph.
y = data.labels
y_hat = model(data.blocks, x)
# Compute loss.
loss = F.cross_entropy(y_hat, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
t1 = time.time()
# Evaluate the model.
acc = evaluate(args, model, graph, features, valid_set, num_classes)
print(
f"Epoch {epoch:05d} | Loss {total_loss / (step + 1):.4f} | "
f"Accuracy {acc.item():.4f} | Time {t1 - t0:.4f}"
)
def parse_args():
parser = argparse.ArgumentParser(
description="A script trains and tests a GraphSAGE model "
"for node classification using GraphBolt dataloader."
)
parser.add_argument(
"--epochs", type=int, default=10, help="Number of training epochs."
)
parser.add_argument(
"--lr",
type=float,
default=1e-3,
help="Learning rate for optimization.",
)
parser.add_argument(
"--batch-size", type=int, default=1024, help="Batch size for training."
)
parser.add_argument(
"--num-workers",
type=int,
default=0,
help="Number of workers for data loading.",
)
parser.add_argument(
"--fanout",
type=str,
default="10,10,10",
help="Fan-out of neighbor sampling. It is IMPORTANT to keep len(fanout)"
" identical with the number of layers in your model. Default: 10,10,10",
)
parser.add_argument(
"--dataset",
type=str,
default="ogbn-products",
choices=[
"ogbn-arxiv",
"ogbn-products",
"ogbn-papers100M",
"igb-hom-tiny",
"igb-hom-small",
"igb-hom-medium",
"igb-hom-large",
"igb-hom",
],
help="The dataset we can use for node classification example. Currently"
" ogbn-products, ogbn-arxiv, ogbn-papers100M and"
" igb-hom-[tiny|small|medium|large] and igb-hom datasets are supported.",
)
parser.add_argument(
"--mode",
default="pinned-cuda",
choices=["cpu-cpu", "cpu-cuda", "pinned-cuda", "cuda-cuda"],
help="Dataset storage placement and Train device: 'cpu' for CPU and RAM,"
" 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
)
parser.add_argument(
"--sample-mode",
default="sample_neighbor",
choices=["sample_neighbor", "sample_layer_neighbor"],
help="The sampling function when doing layerwise sampling.",
)
return parser.parse_args()
def main(args):
if not torch.cuda.is_available():
args.mode = "cpu-cpu"
print(f"Training in {args.mode} mode.")
args.storage_device, args.device = args.mode.split("-")
args.device = torch.device(args.device)
# Load and preprocess dataset.
print("Loading data...")
dataset = gb.BuiltinDataset(args.dataset).load()
# Move the dataset to the selected storage.
if args.storage_device == "pinned":
graph = dataset.graph.pin_memory_()
features = dataset.feature.pin_memory_()
else:
graph = dataset.graph.to(args.storage_device)
features = dataset.feature.to(args.storage_device)
train_set = dataset.tasks[0].train_set
valid_set = dataset.tasks[0].validation_set
test_set = dataset.tasks[0].test_set
all_nodes_set = dataset.all_nodes_set
args.fanout = list(map(int, args.fanout.split(",")))
num_classes = dataset.tasks[0].metadata["num_classes"]
in_size = features.size("node", None, "feat")[0]
hidden_size = 256
out_size = num_classes
model = SAGE(in_size, hidden_size, out_size)
assert len(args.fanout) == len(model.layers)
model = model.to(args.device)
# Model training.
print("Training...")
train(args, graph, features, train_set, valid_set, num_classes, model)
# Test the model.
print("Testing...")
test_acc = layerwise_infer(
args,
graph,
features,
test_set,
all_nodes_set,
model,
num_classes,
)
print(f"Test accuracy {test_acc.item():.4f}")
if __name__ == "__main__":
args = parse_args()
main(args)
examples/graphbolt/pyg/README.md
+57
View File
@@ -0,0 +1,57 @@
## Overview
This project demonstrates the training and evaluation of a GraphSAGE model for node classification on large graphs. The example utilizes GraphBolt for efficient data handling and PyG for the GNN training.
# Node classification on graph
This example aims to demonstrate how to run node classification task on heterogeneous graph with **GraphBolt**.
## Model
The model is a three-layer GraphSAGE network implemented using PyTorch Geometric's SAGEConv layers.
## Default Run on `ogbn-arxiv` dataset
```
python node_classification.py
```
## Accuracies
```
Final performance(for ogbn-arxiv):
All runs:
Highest Train: 62.26
Highest Valid: 59.89
Final Train: 62.26
Final Test: 52.78
```
## Run on `ogbn-products` dataset
### Sample on CPU and train/infer on CPU
```
python node_classification.py --dataset ogbn-products
```
## Accuracies
```
Final performance(for ogbn-products):
All runs:
Highest Train: 90.79
Highest Valid: 89.86
Final Train: 90.79
Final Test: 75.24
```
examples/graphbolt/pyg/hetero/node_classification.py
+529
View File
@@ -0,0 +1,529 @@
"""
This script is a PyG counterpart of ``/examples/graphbolt/rgcn/hetero_rgcn.py``.
"""
import argparse
import time
import dgl.graphbolt as gb
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import SimpleConv
from tqdm import tqdm
def accuracy(out, labels):
assert out.ndim == 2
assert out.size(0) == labels.size(0)
assert labels.ndim == 1 or (labels.ndim == 2 and labels.size(1) == 1)
labels = labels.flatten()
predictions = torch.argmax(out, 1)
return (labels == predictions).sum(dtype=torch.float64) / labels.size(0)
def create_dataloader(
graph,
features,
itemset,
batch_size,
fanout,
device,
job,
):
"""Create a GraphBolt dataloader for training, validation or testing."""
datapipe = gb.ItemSampler(
itemset,
batch_size=batch_size,
shuffle=(job == "train"),
drop_last=(job == "train"),
)
need_copy = True
# Copy the data to the specified device.
if args.graph_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
# Sample neighbors for each node in the mini-batch.
datapipe = getattr(datapipe, args.sample_mode)(
graph,
fanout if job != "infer" else [-1],
overlap_fetch=args.overlap_graph_fetch,
num_gpu_cached_edges=args.num_gpu_cached_edges,
gpu_cache_threshold=args.gpu_graph_caching_threshold,
asynchronous=args.graph_device != "cpu",
)
# Copy the data to the specified device.
if args.feature_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
node_feature_keys = {"paper": ["feat"], "author": ["feat"]}
if args.dataset == "ogb-lsc-mag240m":
node_feature_keys["institution"] = ["feat"]
if "igb-het" in args.dataset:
node_feature_keys["institute"] = ["feat"]
node_feature_keys["fos"] = ["feat"]
# Fetch node features for the sampled subgraph.
datapipe = datapipe.fetch_feature(
features,
node_feature_keys,
overlap_fetch=args.overlap_feature_fetch,
)
# Copy the data to the specified device.
if need_copy:
datapipe = datapipe.copy_to(device=device)
# Create and return a DataLoader to handle data loading.
return gb.DataLoader(datapipe, num_workers=args.num_workers)
class RelGraphConvLayer(nn.Module):
def __init__(
self,
in_size,
out_size,
ntypes,
etypes,
activation,
dropout=0.0,
):
super().__init__()
self.in_size = in_size
self.out_size = out_size
self.activation = activation
# Create a separate convolution layer for each relationship. PyG's
# SimpleConv does not have any weights and only performs message passing
# and aggregation.
self.convs = nn.ModuleDict(
{etype: SimpleConv(aggr="mean") for etype in etypes}
)
# 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(
{
etype: nn.Linear(in_size, out_size, bias=False)
for etype in etypes
}
)
# 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 ntypes}
)
self.dropout = nn.Dropout(dropout)
def forward(self, subgraph, x):
# 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.
(h, h_dst), edge_index, size = subgraph.to_pyg(x)
h_out = {}
for etype in edge_index:
src_ntype, _, dst_ntype = gb.etype_str_to_tuple(etype)
# h_dst is unused in SimpleConv.
t = self.convs[etype](
(h[src_ntype], h_dst[dst_ntype]),
edge_index[etype],
size=size[etype],
)
t = self.weight[etype](t)
if dst_ntype in h_out:
h_out[dst_ntype] += t
else:
h_out[dst_ntype] = t
def _apply(ntype, x):
# 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.
x = x + self.loop_weights[ntype](h_dst[ntype])
return self.dropout(self.activation(x))
# 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 h_out.items()}
class EntityClassify(nn.Module):
def __init__(self, graph, in_size, hidden_size, out_size, n_layers):
super(EntityClassify, self).__init__()
self.layers = nn.ModuleList()
sizes = [in_size] + [hidden_size] * (n_layers - 1) + [out_size]
for i in range(n_layers):
self.layers.append(
RelGraphConvLayer(
sizes[i],
sizes[i + 1],
graph.node_type_to_id.keys(),
graph.edge_type_to_id.keys(),
activation=F.relu if i != n_layers - 1 else lambda x: x,
dropout=0.5,
)
)
def forward(self, subgraphs, h):
for layer, subgraph in zip(self.layers, subgraphs):
h = layer(subgraph, h)
return h
@torch.compile
def evaluate_step(minibatch, model):
category = "paper"
node_features = {
ntype: feat.float()
for (ntype, name), feat in minibatch.node_features.items()
if name == "feat"
}
labels = minibatch.labels[category].long()
out = model(minibatch.sampled_subgraphs, node_features)[category]
num_correct = accuracy(out, labels) * labels.size(0)
return num_correct, labels.size(0)
@torch.no_grad()
def evaluate(
model,
dataloader,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
):
model.eval()
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0
dataloader = tqdm(dataloader, desc="Evaluating")
for step, minibatch in enumerate(dataloader):
num_correct, num_samples = evaluate_step(minibatch, model)
total_correct += num_correct
total_samples += num_samples
if step % 15 == 0:
num_nodes = sum(id.size(0) for id in minibatch.node_ids().values())
dataloader.set_postfix(
{
"num_nodes": num_nodes,
"gpu_cache_miss": gpu_cache_miss_rate_fn(),
"cpu_cache_miss": cpu_cache_miss_rate_fn(),
}
)
return total_correct / total_samples
@torch.compile
def train_step(minibatch, optimizer, model, loss_fn):
category = "paper"
node_features = {
ntype: feat.float()
for (ntype, name), feat in minibatch.node_features.items()
if name == "feat"
}
labels = minibatch.labels[category].long()
optimizer.zero_grad()
out = model(minibatch.sampled_subgraphs, node_features)[category]
loss = loss_fn(out, labels)
# https://github.com/pytorch/pytorch/issues/133942
# num_correct = accuracy(out, labels) * labels.size(0)
num_correct = torch.zeros(1, dtype=torch.float64, device=out.device)
loss.backward()
optimizer.step()
return loss.detach(), num_correct, labels.size(0)
def train_helper(
dataloader,
model,
optimizer,
loss_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
):
model.train()
total_loss = torch.zeros(1, device=device)
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0
start = time.time()
dataloader = tqdm(dataloader, "Training")
for step, minibatch in enumerate(dataloader):
loss, num_correct, num_samples = train_step(
minibatch, optimizer, model, loss_fn
)
total_loss += loss * num_samples
total_correct += num_correct
total_samples += num_samples
if step % 15 == 0:
# log every 15 steps for performance.
num_nodes = sum(id.size(0) for id in minibatch.node_ids().values())
dataloader.set_postfix(
{
"num_nodes": num_nodes,
"gpu_cache_miss": gpu_cache_miss_rate_fn(),
"cpu_cache_miss": cpu_cache_miss_rate_fn(),
}
)
loss = total_loss / total_samples
acc = total_correct / total_samples
end = time.time()
return loss, acc, end - start
def train(
train_dataloader,
valid_dataloader,
model,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
):
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(args.epochs):
train_loss, train_acc, duration = train_helper(
train_dataloader,
model,
optimizer,
loss_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
)
val_acc = evaluate(
model,
valid_dataloader,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
)
print(
f"Epoch: {epoch:02d}, Loss: {train_loss.item():.4f}, "
f"Approx. Train: {train_acc.item():.4f}, "
f"Approx. Val: {val_acc.item():.4f}, "
f"Time: {duration}s"
)
def parse_args():
parser = argparse.ArgumentParser(description="GraphBolt PyG R-SAGE")
parser.add_argument(
"--epochs", type=int, default=10, help="Number of training epochs."
)
parser.add_argument(
"--lr",
type=float,
default=0.001,
help="Learning rate for optimization.",
)
parser.add_argument("--num-hidden", type=int, default=1024)
parser.add_argument(
"--batch-size", type=int, default=1024, help="Batch size for training."
)
parser.add_argument("--num_workers", type=int, default=0)
parser.add_argument(
"--dataset",
type=str,
default="ogb-lsc-mag240m",
choices=[
"ogb-lsc-mag240m",
"igb-het-tiny",
"igb-het-small",
"igb-het-medium",
],
help="Dataset name. Possible values: ogb-lsc-mag240m, igb-het-[tiny|small|medium].",
)
parser.add_argument(
"--fanout",
type=str,
default="25,10",
help="Fan-out of neighbor sampling. It is IMPORTANT to keep len(fanout)"
" identical with the number of layers in your model. Default: 25,10",
)
parser.add_argument(
"--mode",
default="pinned-pinned-cuda",
choices=[
"cpu-cpu-cpu",
"cpu-cpu-cuda",
"cpu-pinned-cuda",
"pinned-pinned-cuda",
"cuda-pinned-cuda",
"cuda-cuda-cuda",
],
help="Graph storage - feature storage - Train device: 'cpu' for CPU and RAM,"
" 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
)
parser.add_argument(
"--sample-mode",
default="sample_neighbor",
choices=["sample_neighbor", "sample_layer_neighbor"],
help="The sampling function when doing layerwise sampling.",
)
parser.add_argument(
"--cpu-feature-cache-policy",
type=str,
default=None,
choices=["s3-fifo", "sieve", "lru", "clock"],
help="The cache policy for the CPU feature cache.",
)
parser.add_argument(
"--cpu-cache-size",
type=float,
default=0,
help="The capacity of the CPU feature cache in GiB.",
)
parser.add_argument(
"--gpu-cache-size",
type=float,
default=0,
help="The capacity of the GPU feature cache in GiB.",
)
parser.add_argument(
"--num-gpu-cached-edges",
type=int,
default=0,
help="The number of edges to be cached from the graph on the GPU.",
)
parser.add_argument(
"--gpu-graph-caching-threshold",
type=int,
default=1,
help="The number of accesses after which a vertex neighborhood will be cached.",
)
parser.add_argument("--precision", type=str, default="high")
return parser.parse_args()
def main():
torch.set_float32_matmul_precision(args.precision)
if not torch.cuda.is_available():
args.mode = "cpu-cpu-cpu"
print(f"Training in {args.mode} mode.")
args.graph_device, args.feature_device, args.device = args.mode.split("-")
args.overlap_feature_fetch = args.feature_device == "pinned"
args.overlap_graph_fetch = args.graph_device == "pinned"
# Load dataset.
dataset = gb.BuiltinDataset(args.dataset).load()
print("Dataset loaded")
# Move the dataset to the selected storage.
graph = (
dataset.graph.pin_memory_()
if args.graph_device == "pinned"
else dataset.graph.to(args.graph_device)
)
features = (
dataset.feature.pin_memory_()
if args.feature_device == "pinned"
else dataset.feature.to(args.feature_device)
)
train_set = dataset.tasks[0].train_set
valid_set = dataset.tasks[0].validation_set
test_set = dataset.tasks[0].test_set
args.fanout = list(map(int, args.fanout.split(",")))
num_classes = dataset.tasks[0].metadata["num_classes"]
num_etypes = len(graph.num_edges)
feats_on_disk = {
k: features[k]
for k in features.keys()
if k[2] == "feat" and isinstance(features[k], gb.DiskBasedFeature)
}
if args.cpu_cache_size > 0 and len(feats_on_disk) > 0:
cached_features = gb.cpu_cached_feature(
feats_on_disk,
int(args.cpu_cache_size * (2**30)),
args.cpu_feature_cache_policy,
args.feature_device == "pinned",
)
for k, cpu_cached_feature in cached_features.items():
features[k] = cpu_cached_feature
cpu_cache_miss_rate_fn = lambda: cpu_cached_feature.miss_rate
else:
cpu_cache_miss_rate_fn = lambda: 1
if args.gpu_cache_size > 0 and args.feature_device != "cuda":
feats = {k: features[k] for k in features.keys() if k[2] == "feat"}
cached_features = gb.gpu_cached_feature(
feats,
int(args.gpu_cache_size * (2**30)),
)
for k, gpu_cached_feature in cached_features.items():
features[k] = gpu_cached_feature
gpu_cache_miss_rate_fn = lambda: gpu_cached_feature.miss_rate
else:
gpu_cache_miss_rate_fn = lambda: 1
train_dataloader, valid_dataloader, test_dataloader = (
create_dataloader(
graph=graph,
features=features,
itemset=itemset,
batch_size=args.batch_size,
fanout=[
torch.full((num_etypes,), fanout) for fanout in args.fanout
],
device=args.device,
job=job,
)
for itemset, job in zip(
[train_set, valid_set, test_set], ["train", "evaluate", "evaluate"]
)
)
feat_size = features.size("node", "paper", "feat")[0]
hidden_channels = args.num_hidden
# Initialize the entity classification model.
model = EntityClassify(
graph, feat_size, hidden_channels, num_classes, len(args.fanout)
).to(args.device)
print(
"Number of model parameters: "
f"{sum(p.numel() for p in model.parameters())}"
)
train(
train_dataloader,
valid_dataloader,
model,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
args.device,
)
# Labels are currently unavailable for mag240M so the test acc will be 0.
print("Testing...")
test_acc = evaluate(
model,
test_dataloader,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
args.device,
)
print(f"Test accuracy {test_acc.item():.4f}")
if __name__ == "__main__":
args = parse_args()
main()
examples/graphbolt/pyg/labor/README.md
+94
View File
@@ -0,0 +1,94 @@
Layer-Neighbor Sampling -- Defusing Neighborhood Explosion in GNNs
============
- Paper link: [https://papers.nips.cc/paper_files/paper/2023/hash/51f9036d5e7ae822da8f6d4adda1fb39-Abstract-Conference.html](NeurIPS 2023)
This is an official Labor sampling example to showcase the use of [https://docs.dgl.ai/en/latest/generated/dgl.graphbolt.LayerNeighborSampler.html](dgl.graphbolt.LayerNeighborSampler).
This sampler has 2 parameters, `layer_dependency=[False|True]` and
`batch_dependency=k`, where k is any nonnegative integer.
We use early stopping so that the final accuracy numbers are reported with a
fairly well converged model. Additional contributions to improve the validation
accuracy are welcome, and hence hopefully also improving the test accuracy.
### layer_dependency
Enabling this parameter by the command line option `--layer-dependency` makes it so
that the random variates for sampling are identical across layers. This ensures
that the same vertex gets the same neighborhood in each layer.
### batch_dependency
This method is proposed in Section 3.2 of [https://arxiv.org/pdf/2310.12403](Cooperative Minibatching in Graph Neural Networks), it is denoted as kappa in the paper. It
makes the random variates used across minibatches dependent, thus increasing
temporal locality. When used with a cache, the increase in the temporal locality
can be observed by monitoring the drop in the cache miss rate with higher values
of the batch dependency parameter, speeding up embedding transfers to the GPU.
### Performance
Use the `--torch-compile` option for best performance. If your GPU has spare
memory, consider using `--mode=cuda-cuda-cuda` to move the whole dataset to the
GPU. If not, consider using `--mode=cuda-pinned-cuda --num-gpu-cached-features=N`
to keep the graph on the GPU and features in system RAM with `N` of the node
features cached on the GPU. If you can not even fit the graph on the GPU, then
consider using `--mode=pinned-pinned-cuda --num-gpu-cached-features=N`. Finally,
you can use `--mode=cpu-pinned=cuda --num-gpu-cached-features=N` to perform the
sampling operation on the CPU.
### Examples
We use `--num-gpu-cached-features=500000` to cache the 500k of the node
embeddings for the `ogbn-products` dataset (default). Check the command line
arguments to see which other datasets can be run. When running with the yelp
dataset, using `--dropout=0` gives better final validation and test accuracy.
Example run with batch_dependency=1, cache miss rate is 62%:
```bash
python node_classification.py --num-gpu-cached-features=500000 --batch-dependency=1
Training in pinned-pinned-cuda mode.
Loading data...
The dataset is already preprocessed.
Training: 192it [00:03, 50.95it/s, num_nodes=247243, cache_miss=0.619]
Evaluating: 39it [00:00, 76.01it/s, num_nodes=137466, cache_miss=0.621]
Epoch 00, Loss: 1.1161, Approx. Train: 0.7024, Approx. Val: 0.8612, Time: 3.7688188552856445s
```
Example run with batch_dependency=32, cache miss rate is 22%:
```bash
python node_classification.py --num-gpu-cached-features=500000 --batch-dependency=32
Training in pinned-pinned-cuda mode.
Loading data...
The dataset is already preprocessed.
Training: 192it [00:03, 54.34it/s, num_nodes=250479, cache_miss=0.221]
Evaluating: 39it [00:00, 84.66it/s, num_nodes=135142, cache_miss=0.226]
Epoch 00, Loss: 1.1288, Approx. Train: 0.6993, Approx. Val: 0.8607, Time: 3.5339605808258057s
```
Example run with layer_dependency=True, # sampled nodes is 190k vs 250k without
this option:
```bash
python node_classification.py --num-gpu-cached-features=500000 --layer-dependency
Training in pinned-pinned-cuda mode.
Loading data...
The dataset is already preprocessed.
Training: 192it [00:03, 54.03it/s, num_nodes=191259, cache_miss=0.626]
Evaluating: 39it [00:00, 79.49it/s, num_nodes=108720, cache_miss=0.627]
Epoch 00, Loss: 1.1495, Approx. Train: 0.6932, Approx. Val: 0.8586, Time: 3.5540308952331543s
```
Example run with the original GraphSAGE sampler (Neighbor Sampler), # sampled nodes
is 520k, more than 2x higher than Labor sampler.
```bash
python node_classification.py --num-gpu-cached-features=500000 --sample-mode=sample_neighbor
Training in pinned-pinned-cuda mode.
Loading data...
The dataset is already preprocessed.
Training: 192it [00:04, 45.60it/s, num_nodes=517522, cache_miss=0.563]
Evaluating: 39it [00:00, 77.53it/s, num_nodes=255686, cache_miss=0.565]
Epoch 00, Loss: 1.1152, Approx. Train: 0.7015, Approx. Val: 0.8652, Time: 4.211000919342041s
```
examples/graphbolt/pyg/labor/load_dataset.py
+55
View File
@@ -0,0 +1,55 @@
import dgl.graphbolt as gb
def load_dgl(name):
from dgl.data import (
CiteseerGraphDataset,
CoraGraphDataset,
FlickrDataset,
PubmedGraphDataset,
RedditDataset,
YelpDataset,
)
d = {
"cora": CoraGraphDataset,
"citeseer": CiteseerGraphDataset,
"pubmed": PubmedGraphDataset,
"reddit": RedditDataset,
"yelp": YelpDataset,
"flickr": FlickrDataset,
}
dataset = gb.LegacyDataset(d[name]())
new_feature = gb.TorchBasedFeatureStore([])
new_feature._features = dataset.feature._features
dataset._feature = new_feature
multilabel = name in ["yelp"]
return dataset, multilabel
def load_dataset(dataset_name, disk_based_feature_keys=None):
multilabel = False
if dataset_name in [
"reddit",
"cora",
"citeseer",
"pubmed",
"yelp",
"flickr",
]:
dataset, multilabel = load_dgl(dataset_name)
else:
if "mag240M" in dataset_name:
dataset_name = "ogb-lsc-mag240m"
dataset = gb.BuiltinDataset(dataset_name)
if disk_based_feature_keys is None:
disk_based_feature_keys = set()
for feature in dataset.yaml_data["feature_data"]:
feature_key = (feature["domain"], feature["type"], feature["name"])
# Set the in_memory setting to False without modifying YAML file.
if feature_key in disk_based_feature_keys:
feature["in_memory"] = False
dataset = dataset.load()
return dataset, multilabel
examples/graphbolt/pyg/labor/node_classification.py
+569
View File
@@ -0,0 +1,569 @@
import argparse
import time
from copy import deepcopy
from functools import partial
import dgl.graphbolt as gb
import torch
# For torch.compile until https://github.com/pytorch/pytorch/issues/121197 is
# resolved.
import torch._inductor.codecache
torch._dynamo.config.cache_size_limit = 32
import torch.nn as nn
import torchmetrics.functional as MF
from load_dataset import load_dataset
from sage_conv import SAGEConv as CustomSAGEConv
from torch_geometric.nn import SAGEConv
from tqdm import tqdm
def accuracy(out, labels):
assert out.ndim == 2
assert out.size(0) == labels.size(0)
assert labels.ndim == 1 or (labels.ndim == 2 and labels.size(1) == 1)
labels = labels.flatten()
predictions = torch.argmax(out, 1)
return (labels == predictions).sum(dtype=torch.float64) / labels.size(0)
class GraphSAGE(torch.nn.Module):
def __init__(
self, in_size, hidden_size, out_size, n_layers, dropout, variant
):
super().__init__()
assert variant in ["original", "custom"]
self.layers = torch.nn.ModuleList()
if variant == "custom":
sizes = [in_size] + [hidden_size] * n_layers
for i in range(n_layers):
self.layers.append(CustomSAGEConv(sizes[i], sizes[i + 1]))
self.linear = nn.Linear(hidden_size, out_size)
self.activation = nn.GELU()
else:
sizes = [in_size] + [hidden_size] * (n_layers - 1) + [out_size]
for i in range(n_layers):
self.layers.append(SAGEConv(sizes[i], sizes[i + 1]))
self.activation = nn.ReLU()
self.dropout = nn.Dropout(dropout)
self.hidden_size = hidden_size
self.out_size = out_size
self.variant = variant
def forward(self, subgraphs, x):
h = x
for i, (layer, subgraph) in enumerate(zip(self.layers, subgraphs)):
h, edge_index, size = subgraph.to_pyg(h)
h = layer(h, edge_index, size=size)
if self.variant == "custom":
h = self.activation(h)
h = self.dropout(h)
elif i != len(subgraphs) - 1:
h = self.activation(h)
return self.linear(h) if self.variant == "custom" else h
def inference(self, graph, features, dataloader, storage_device):
"""Conduct layer-wise inference to get all the node embeddings."""
pin_memory = storage_device == "pinned"
buffer_device = torch.device("cpu" if pin_memory else storage_device)
for layer_idx, layer in enumerate(self.layers):
is_last_layer = layer_idx == len(self.layers) - 1
y = torch.empty(
graph.total_num_nodes,
self.out_size if is_last_layer else self.hidden_size,
dtype=torch.float32,
device=buffer_device,
pin_memory=pin_memory,
)
for data in tqdm(dataloader, "Inferencing"):
# len(data.sampled_subgraphs) = 1
h, edge_index, size = data.sampled_subgraphs[0].to_pyg(
data.node_features["feat"]
)
hidden_x = layer(h, edge_index, size=size)
if self.variant == "custom":
hidden_x = self.activation(hidden_x)
if is_last_layer:
hidden_x = self.linear(hidden_x)
elif not is_last_layer:
hidden_x = self.activation(hidden_x)
# By design, our output nodes are contiguous.
y[data.seeds[0] : data.seeds[-1] + 1] = hidden_x.to(
buffer_device
)
if not is_last_layer:
features.update("node", None, "feat", y)
return y
def create_dataloader(
graph, features, itemset, batch_size, fanout, device, job
):
# Initialize an ItemSampler to sample mini-batches from the dataset.
datapipe = gb.ItemSampler(
itemset,
batch_size=batch_size,
shuffle=(job == "train"),
drop_last=(job == "train"),
)
need_copy = True
# Copy the data to the specified device.
if args.graph_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
# Sample neighbors for each node in the mini-batch.
kwargs = (
{
"layer_dependency": args.layer_dependency,
"batch_dependency": args.batch_dependency,
}
if args.sample_mode == "sample_layer_neighbor"
else {}
)
datapipe = getattr(datapipe, args.sample_mode)(
graph,
fanout if job != "infer" else [-1],
overlap_fetch=args.overlap_graph_fetch,
asynchronous=args.graph_device != "cpu",
**kwargs,
)
# Copy the data to the specified device.
if args.feature_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
# Fetch node features for the sampled subgraph.
datapipe = datapipe.fetch_feature(
features,
node_feature_keys=["feat"],
overlap_fetch=args.overlap_feature_fetch,
)
# Copy the data to the specified device.
if need_copy:
datapipe = datapipe.copy_to(device=device)
# Create and return a DataLoader to handle data loading.
return gb.DataLoader(datapipe, num_workers=args.num_workers)
@torch.compile
def train_step(minibatch, optimizer, model, loss_fn, multilabel, eval_fn):
node_features = minibatch.node_features["feat"]
labels = minibatch.labels
optimizer.zero_grad()
out = model(minibatch.sampled_subgraphs, node_features)
label_dtype = out.dtype if multilabel else None
loss = loss_fn(out, labels.to(label_dtype))
num_correct = eval_fn(out, labels) * labels.size(0)
loss.backward()
optimizer.step()
return loss.detach(), num_correct, labels.size(0)
def train_helper(
dataloader,
model,
optimizer,
loss_fn,
multilabel,
eval_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
):
model.train() # Set the model to training mode
total_loss = torch.zeros(1, device=device) # Accumulator for the total loss
# Accumulator for the total number of correct predictions
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0 # Accumulator for the total number of samples processed
num_batches = 0 # Counter for the number of mini-batches processed
start = time.time()
dataloader = tqdm(dataloader, "Training")
for step, minibatch in enumerate(dataloader):
loss, num_correct, num_samples = train_step(
minibatch, optimizer, model, loss_fn, multilabel, eval_fn
)
total_loss += loss
total_correct += num_correct
total_samples += num_samples
num_batches += 1
if step % 25 == 0:
# log every 25 steps for performance.
dataloader.set_postfix(
{
"num_nodes": minibatch.node_ids().size(0),
"gpu_cache_miss": gpu_cache_miss_rate_fn(),
"cpu_cache_miss": cpu_cache_miss_rate_fn(),
}
)
train_loss = total_loss / num_batches
train_acc = total_correct / total_samples
end = time.time()
return train_loss, train_acc, end - start
def train(
train_dataloader,
valid_dataloader,
model,
multilabel,
eval_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
):
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
loss_fn = nn.BCEWithLogitsLoss() if multilabel else nn.CrossEntropyLoss()
best_model = None
best_model_acc = 0
best_model_epoch = -1
for epoch in range(args.epochs):
train_loss, train_acc, duration = train_helper(
train_dataloader,
model,
optimizer,
loss_fn,
multilabel,
eval_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
)
val_acc = evaluate(
model,
valid_dataloader,
eval_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
)
if val_acc > best_model_acc:
best_model_acc = val_acc
best_model = deepcopy(model.state_dict())
best_model_epoch = epoch
print(
f"Epoch {epoch:02d}, Loss: {train_loss.item():.4f}, "
f"Approx. Train: {train_acc.item():.4f}, "
f"Approx. Val: {val_acc.item():.4f}, "
f"Time: {duration}s"
)
if best_model_epoch + args.early_stopping_patience < epoch:
break
return best_model
@torch.no_grad()
def layerwise_infer(
args,
graph,
features,
itemsets,
all_nodes_set,
model,
eval_fn,
):
model.eval()
dataloader = create_dataloader(
graph=graph,
features=features,
itemset=all_nodes_set,
batch_size=args.batch_size,
fanout=[-1],
device=args.device,
job="infer",
)
pred = model.inference(graph, features, dataloader, args.feature_device)
metrics = {}
for split_name, itemset in itemsets.items():
nid, labels = itemset[:]
acc = eval_fn(
pred[nid.to(pred.device)],
labels.to(pred.device),
)
metrics[split_name] = acc.item()
return metrics
@torch.compile
def evaluate_step(minibatch, model, eval_fn):
node_features = minibatch.node_features["feat"]
labels = minibatch.labels
out = model(minibatch.sampled_subgraphs, node_features)
num_correct = eval_fn(out, labels) * labels.size(0)
return num_correct, labels.size(0)
@torch.no_grad()
def evaluate(
model,
dataloader,
eval_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
device,
):
model.eval()
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0
dataloader = tqdm(dataloader, "Evaluating")
for step, minibatch in enumerate(dataloader):
num_correct, num_samples = evaluate_step(minibatch, model, eval_fn)
total_correct += num_correct
total_samples += num_samples
if step % 25 == 0:
dataloader.set_postfix(
{
"num_nodes": minibatch.node_ids().size(0),
"gpu_cache_miss": gpu_cache_miss_rate_fn(),
"cpu_cache_miss": cpu_cache_miss_rate_fn(),
}
)
return total_correct / total_samples
def parse_args():
parser = argparse.ArgumentParser(
description="Which dataset are you going to use?"
)
parser.add_argument(
"--epochs", type=int, default=9999999, help="Number of training epochs."
)
parser.add_argument(
"--lr",
type=float,
default=0.001,
help="Learning rate for optimization.",
)
parser.add_argument("--num-hidden", type=int, default=256)
parser.add_argument("--dropout", type=float, default=0.5)
parser.add_argument(
"--batch-size", type=int, default=1024, help="Batch size for training."
)
parser.add_argument(
"--num-workers",
type=int,
default=0,
help="Number of workers for data loading.",
)
parser.add_argument(
"--dataset",
type=str,
default="ogbn-products",
choices=[
"ogbn-arxiv",
"ogbn-products",
"ogbn-papers100M",
"igb-hom-tiny",
"igb-hom-small",
"igb-hom-medium",
"igb-hom-large",
"igb-hom",
"reddit",
"yelp",
"flickr",
],
)
parser.add_argument(
"--fanout",
type=str,
default="10,10,10",
help="Fan-out of neighbor sampling. len(fanout) determines the number of"
" GNN layers in your model. Default: 10,10,10",
)
parser.add_argument(
"--mode",
default="pinned-pinned-cuda",
choices=[
"cpu-cpu-cpu",
"cpu-cpu-cuda",
"cpu-pinned-cuda",
"pinned-pinned-cuda",
"cuda-pinned-cuda",
"cuda-cuda-cuda",
],
help="Graph storage - feature storage - Train device: 'cpu' for CPU and"
" RAM, 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
)
parser.add_argument("--layer-dependency", action="store_true")
parser.add_argument("--batch-dependency", type=int, default=1)
parser.add_argument(
"--cpu-feature-cache-policy",
type=str,
default=None,
choices=["s3-fifo", "sieve", "lru", "clock"],
help="The cache policy for the CPU feature cache.",
)
parser.add_argument(
"--num-cpu-cached-features",
type=int,
default=0,
help="The capacity of the CPU cache, the number of features to store.",
)
parser.add_argument(
"--num-gpu-cached-features",
type=int,
default=0,
help="The capacity of the GPU cache, the number of features to store.",
)
parser.add_argument("--early-stopping-patience", type=int, default=25)
parser.add_argument(
"--sample-mode",
default="sample_layer_neighbor",
choices=["sample_neighbor", "sample_layer_neighbor"],
help="The sampling function when doing layerwise sampling.",
)
parser.add_argument(
"--sage-model-variant",
default="custom",
choices=["custom", "original"],
help="The custom SAGE GNN model provides higher accuracy with lower"
" runtime performance.",
)
parser.add_argument("--precision", type=str, default="high")
return parser.parse_args()
def main():
torch.set_float32_matmul_precision(args.precision)
if not torch.cuda.is_available():
args.mode = "cpu-cpu-cpu"
print(f"Training in {args.mode} mode.")
args.graph_device, args.feature_device, args.device = args.mode.split("-")
args.overlap_feature_fetch = args.feature_device == "pinned"
args.overlap_graph_fetch = args.graph_device == "pinned"
# Load and preprocess dataset.
print("Loading data...")
disk_based_feature_keys = None
if args.num_cpu_cached_features > 0:
disk_based_feature_keys = [("node", None, "feat")]
dataset, multilabel = load_dataset(args.dataset, disk_based_feature_keys)
# Move the dataset to the selected storage.
graph = (
dataset.graph.pin_memory_()
if args.graph_device == "pinned"
else dataset.graph.to(args.graph_device)
)
features = (
dataset.feature.pin_memory_()
if args.feature_device == "pinned"
else dataset.feature.to(args.feature_device)
)
train_set = dataset.tasks[0].train_set
valid_set = dataset.tasks[0].validation_set
test_set = dataset.tasks[0].test_set
all_nodes_set = dataset.all_nodes_set
args.fanout = list(map(int, args.fanout.split(",")))
num_classes = dataset.tasks[0].metadata["num_classes"]
feature_index_device = (
args.feature_device if args.feature_device != "pinned" else None
)
feature_num_bytes = (
features[("node", None, "feat")]
# Read a single row to query its size in bytes.
.read(torch.zeros(1, device=feature_index_device).long()).nbytes
)
if args.num_cpu_cached_features > 0 and isinstance(
features[("node", None, "feat")], gb.DiskBasedFeature
):
features[("node", None, "feat")] = gb.cpu_cached_feature(
features[("node", None, "feat")],
args.num_cpu_cached_features * feature_num_bytes,
args.cpu_feature_cache_policy,
args.feature_device == "pinned",
)
cpu_cached_feature = features[("node", None, "feat")]
cpu_cache_miss_rate_fn = lambda: cpu_cached_feature.miss_rate
else:
cpu_cache_miss_rate_fn = lambda: 1
if args.num_gpu_cached_features > 0 and args.feature_device != "cuda":
features[("node", None, "feat")] = gb.gpu_cached_feature(
features[("node", None, "feat")],
args.num_gpu_cached_features * feature_num_bytes,
)
gpu_cached_feature = features[("node", None, "feat")]
gpu_cache_miss_rate_fn = lambda: gpu_cached_feature.miss_rate
else:
gpu_cache_miss_rate_fn = lambda: 1
train_dataloader, valid_dataloader = (
create_dataloader(
graph=graph,
features=features,
itemset=itemset,
batch_size=args.batch_size,
fanout=args.fanout,
device=args.device,
job=job,
)
for itemset, job in zip([train_set, valid_set], ["train", "evaluate"])
)
in_channels = features.size("node", None, "feat")[0]
model = GraphSAGE(
in_channels,
args.num_hidden,
num_classes,
len(args.fanout),
args.dropout,
args.sage_model_variant,
).to(args.device)
assert len(args.fanout) == len(model.layers)
eval_fn = (
partial(
# TODO @mfbalin: Find an implementation that does not synchronize.
MF.f1_score,
task="multilabel",
num_labels=num_classes,
validate_args=False,
)
if multilabel
else accuracy
)
best_model = train(
train_dataloader,
valid_dataloader,
model,
multilabel,
eval_fn,
gpu_cache_miss_rate_fn,
cpu_cache_miss_rate_fn,
args.device,
)
model.load_state_dict(best_model)
# Test the model.
print("Testing...")
itemsets = {"train": train_set, "val": valid_set, "test": test_set}
final_acc = layerwise_infer(
args,
graph,
features,
itemsets,
all_nodes_set,
model,
eval_fn,
)
print("Final accuracy values:")
print(final_acc)
if __name__ == "__main__":
args = parse_args()
main()
examples/graphbolt/pyg/labor/sage_conv.py
+145
View File
@@ -0,0 +1,145 @@
from typing import List, Optional, Tuple, Union
import torch.nn.functional as F
from torch import Tensor
from torch_geometric.nn.aggr import Aggregation, MultiAggregation
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.nn.dense.linear import Linear
from torch_geometric.typing import Adj, OptPairTensor, Size, SparseTensor
from torch_geometric.utils import spmm
class SAGEConv(MessagePassing):
r"""A variant of the GraphSAGE operator from the `"Inductive Representation
Learning on Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper.
.. math::
\mathbf{x}^{\prime}_i = \mathbf{W}_1 \mathbf{x}_i + \mathbf{W}_2 \cdot
\mathrm{mean}_{j \in \mathcal{N(i)}} \mathbf{x}_j
If :obj:`project = True`, then :math:`\mathbf{x}_j` will first get
projected via
.. math::
\mathbf{x}_j \leftarrow \sigma ( \mathbf{W}_3 \mathbf{x}_j +
\mathbf{b})
as described in Eq. (3) of the paper.
Args:
in_channels (int or tuple): Size of each input sample, or :obj:`-1` to
derive the size from the first input(s) to the forward method.
A tuple corresponds to the sizes of source and target
dimensionalities.
out_channels (int): Size of each output sample.
aggr (str or Aggregation, optional): The aggregation scheme to use.
Any aggregation of :obj:`torch_geometric.nn.aggr` can be used,
*e.g.*, :obj:`"mean"`, :obj:`"max"`, or :obj:`"lstm"`.
(default: :obj:`"mean"`)
project (bool, optional): If set to :obj:`True`, the layer will apply a
linear transformation followed by an activation function before
aggregation (as described in Eq. (3) of the paper).
(default: :obj:`True`)
bias (bool, optional): If set to :obj:`False`, the layer will not learn
an additive bias. (default: :obj:`True`)
**kwargs (optional): Additional arguments of
:class:`torch_geometric.nn.conv.MessagePassing`.
Shapes:
- **inputs:**
node features :math:`(|\mathcal{V}|, F_{in})` or
:math:`((|\mathcal{V_s}|, F_{s}), (|\mathcal{V_t}|, F_{t}))`
if bipartite,
edge indices :math:`(2, |\mathcal{E}|)`
- **outputs:** node features :math:`(|\mathcal{V}|, F_{out})` or
:math:`(|\mathcal{V_t}|, F_{out})` if bipartite
"""
def __init__(
self,
in_channels: Union[int, Tuple[int, int]],
out_channels: int,
aggr: Optional[Union[str, List[str], Aggregation]] = "mean",
project: bool = True,
bias: bool = True,
**kwargs,
):
self.in_channels = in_channels
self.out_channels = out_channels
self.project = project
if isinstance(in_channels, int):
in_channels = (in_channels, in_channels)
if aggr == "lstm":
kwargs.setdefault("aggr_kwargs", {})
kwargs["aggr_kwargs"].setdefault("in_channels", in_channels[0])
kwargs["aggr_kwargs"].setdefault("out_channels", in_channels[0])
super().__init__(aggr, **kwargs)
if self.project:
if in_channels[0] <= 0:
raise ValueError(
f"'{self.__class__.__name__}' does not "
f"support lazy initialization with "
f"`project=True`"
)
self.lin = Linear(in_channels[0], in_channels[0], bias=True)
if isinstance(self.aggr_module, MultiAggregation):
aggr_out_channels = self.aggr_module.get_out_channels(
in_channels[0]
)
else:
aggr_out_channels = in_channels[0]
self.lin_l = Linear(aggr_out_channels, out_channels, bias=bias)
self.lin_r = Linear(in_channels[1], out_channels, bias=False)
self.reset_parameters()
def reset_parameters(self):
super().reset_parameters()
if self.project:
self.lin.reset_parameters()
self.lin_l.reset_parameters()
self.lin_r.reset_parameters()
def forward(
self,
x: Union[Tensor, OptPairTensor],
edge_index: Adj,
size: Size = None,
) -> Tensor:
if isinstance(x, Tensor):
x = (x, x)
if self.project and hasattr(self, "lin"):
x = (F.gelu(self.lin(x[0])), x[1])
# propagate_type: (x: OptPairTensor)
AX = self.propagate(edge_index, x=x, size=size)
out = self.lin_l(AX)
x_r = x[1]
if x_r is not None:
out = out + self.lin_r(x_r)
return out
def message(self, x_j: Tensor) -> Tensor:
return x_j
def message_and_aggregate(self, adj_t: Adj, x: OptPairTensor) -> Tensor:
if isinstance(adj_t, SparseTensor):
adj_t = adj_t.set_value(None, layout=None)
return spmm(adj_t, x[0], reduce=self.aggr)
def __repr__(self) -> str:
return (
f"{self.__class__.__name__}({self.in_channels}, "
f"{self.out_channels}, aggr={self.aggr})"
)
examples/graphbolt/pyg/link_prediction.py
+463
View File
@@ -0,0 +1,463 @@
"""
This script trains and tests a GraphSAGE model for link prediction on
large graphs using graphbolt dataloader. It is the PyG counterpart of the
example in `examples/graphbolt/link_prediction.py`.
Paper: [Inductive Representation Learning on Large Graphs]
(https://arxiv.org/abs/1706.02216)
While node classification predicts labels for nodes based on their
local neighborhoods, link prediction assesses the likelihood of an edge
existing between two nodes, necessitating different sampling strategies
that account for pairs of nodes and their joint neighborhoods.
This flowchart describes the main functional sequence of the provided example.
main
│
├───> OnDiskDataset pre-processing
│
├───> Instantiate SAGE model
│
├───> train
│ │
│ ├───> Get graphbolt dataloader (HIGHLIGHT)
| |
| |───> Define a PyG GNN model for link prediction (HIGHLIGHT)
│ │
│ └───> Training loop
│ │
│ ├───> SAGE.forward
│
└───> Validation and test set evaluation
"""
import argparse
import time
from functools import partial
import dgl.graphbolt as gb
import torch
# For torch.compile until https://github.com/pytorch/pytorch/issues/121197 is
# resolved.
import torch._inductor.codecache
torch._dynamo.config.cache_size_limit = 32
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from torchmetrics.retrieval import RetrievalMRR
from tqdm import tqdm, trange
class GraphSAGE(torch.nn.Module):
#####################################################################
# (HIGHLIGHT) Define the GraphSAGE model architecture.
#
# - This class inherits from `torch.nn.Module`.
# - Two convolutional layers are created using the SAGEConv class from PyG.
# - The forward method defines the computation performed at every call.
#####################################################################
def __init__(self, in_size, hidden_size, n_layers):
super(GraphSAGE, self).__init__()
self.layers = torch.nn.ModuleList()
sizes = [in_size] + [hidden_size] * n_layers
for i in range(n_layers):
self.layers.append(SAGEConv(sizes[i], sizes[i + 1]))
self.hidden_size = hidden_size
self.predictor = torch.nn.Sequential(
torch.nn.Linear(hidden_size, hidden_size),
torch.nn.ReLU(),
torch.nn.Linear(hidden_size, hidden_size),
torch.nn.ReLU(),
torch.nn.Linear(hidden_size, 1),
)
def forward(self, subgraphs, x):
h = x
for i, (layer, subgraph) in enumerate(zip(self.layers, subgraphs)):
#####################################################################
# (HIGHLIGHT) Convert given features to be consumed by a PyG layer.
#
# PyG layers have two modes, bipartite and normal. We slice the
# given features to get src and dst features to use the PyG layers
# in the more efficient bipartite mode.
#####################################################################
h, edge_index, size = subgraph.to_pyg(h)
h = layer(h, edge_index, size=size)
if i != len(subgraphs) - 1:
h = F.relu(h)
return h
def inference(self, graph, features, dataloader, storage_device):
"""Conduct layer-wise inference to get all the node embeddings."""
pin_memory = storage_device == "pinned"
buffer_device = torch.device("cpu" if pin_memory else storage_device)
for layer_idx, layer in enumerate(self.layers):
is_last_layer = layer_idx == len(self.layers) - 1
y = torch.empty(
graph.total_num_nodes,
self.hidden_size,
dtype=torch.float32,
device=buffer_device,
pin_memory=pin_memory,
)
for data in tqdm(dataloader, "Inferencing"):
# len(data.sampled_subgraphs) = 1
h, edge_index, size = data.sampled_subgraphs[0].to_pyg(
data.node_features["feat"]
)
hidden_x = layer(h, edge_index, size=size)
if not is_last_layer:
hidden_x = F.relu(hidden_x)
# By design, our output nodes are contiguous.
y[data.seeds[0] : data.seeds[-1] + 1] = hidden_x.to(
buffer_device
)
if not is_last_layer:
features.update("node", None, "feat", y)
return y
def create_dataloader(
graph, features, itemset, batch_size, fanout, device, job
):
#####################################################################
# (HIGHLIGHT) Create a data loader for efficiently loading graph data.
#
# - 'ItemSampler' samples mini-batches of node IDs from the dataset.
# - 'CopyTo' copies the fetched data to the specified device.
# - 'sample_neighbor' performs neighbor sampling on the graph.
# - 'FeatureFetcher' fetches node features based on the sampled subgraph.
#####################################################################
# Create a datapipe for mini-batch sampling with a specific neighbor fanout.
# Here, [10, 10, 10] specifies the number of neighbors sampled for each node at each layer.
# We're using `sample_neighbor` for consistency with DGL's sampling API.
# Note: GraphBolt offers additional sampling methods, such as `sample_layer_neighbor`,
# which could provide further optimization and efficiency for GNN training.
# Users are encouraged to explore these advanced features for potentially improved performance.
# Initialize an ItemSampler to sample mini-batches from the dataset.
datapipe = gb.ItemSampler(
itemset,
batch_size=batch_size,
shuffle=(job == "train"),
drop_last=(job == "train"),
)
need_copy = True
# Copy the data to the specified device.
if args.graph_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
# Sample negative edges.
if job == "train":
datapipe = datapipe.sample_uniform_negative(graph, args.neg_ratio)
# Sample neighbors for each node in the mini-batch.
datapipe = getattr(datapipe, args.sample_mode)(
graph,
fanout if job != "infer" else [-1],
overlap_fetch=args.overlap_graph_fetch,
asynchronous=args.graph_device != "cpu",
)
if job == "train" and args.exclude_edges:
datapipe = datapipe.exclude_seed_edges(
include_reverse_edges=True,
asynchronous=args.graph_device != "cpu",
)
# Copy the data to the specified device.
if args.feature_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
# Fetch node features for the sampled subgraph.
datapipe = datapipe.fetch_feature(
features,
node_feature_keys=["feat"],
overlap_fetch=args.overlap_feature_fetch,
)
# Copy the data to the specified device.
if need_copy:
datapipe = datapipe.copy_to(device=device)
# Create and return a DataLoader to handle data loading.
return gb.DataLoader(datapipe, num_workers=args.num_workers)
@torch.compile
def predictions_step(model, h_src, h_dst):
return model.predictor(h_src * h_dst).squeeze()
def compute_predictions(model, node_emb, seeds, device):
"""Compute the predictions for given source and destination nodes.
This function computes the predictions for a set of node pairs, dividing the
task into batches to handle potentially large graphs.
"""
preds = torch.empty(seeds.shape[0], device=device)
seeds_src, seeds_dst = seeds.T
# The constant number is 1001, due to negtive ratio in the `ogbl-citation2`
# dataset is 1000.
eval_size = args.eval_batch_size * 1001
# Loop over node pairs in batches.
for start in trange(0, seeds_src.shape[0], eval_size, desc="Evaluate"):
end = min(start + eval_size, seeds_src.shape[0])
# Fetch embeddings for current batch of source and destination nodes.
h_src = node_emb[seeds_src[start:end]].to(device, non_blocking=True)
h_dst = node_emb[seeds_dst[start:end]].to(device, non_blocking=True)
# Compute prediction scores using the model.
preds[start:end] = predictions_step(model, h_src, h_dst)
return preds
@torch.no_grad()
def evaluate(model, graph, features, all_nodes_set, valid_set, test_set):
"""Evaluate the model on validation and test sets."""
model.eval()
dataloader = create_dataloader(
graph,
features,
all_nodes_set,
args.eval_batch_size,
[-1],
args.device,
job="infer",
)
# Compute node embeddings for the entire graph.
node_emb = model.inference(graph, features, dataloader, args.feature_device)
results = []
# Loop over both validation and test sets.
for split in [valid_set, test_set]:
# Unpack the item set.
seeds = split._items[0].to(node_emb.device)
labels = split._items[1].to(node_emb.device)
indexes = split._items[2].to(node_emb.device)
preds = compute_predictions(model, node_emb, seeds, indexes.device)
# Compute MRR values for the current split.
results.append(RetrievalMRR()(preds, labels, indexes))
return results
@torch.compile
def train_step(minibatch, optimizer, model):
node_features = minibatch.node_features["feat"]
compacted_seeds = minibatch.compacted_seeds.T
labels = minibatch.labels
optimizer.zero_grad()
y = model(minibatch.sampled_subgraphs, node_features)
logits = model.predictor(
y[compacted_seeds[0]] * y[compacted_seeds[1]]
).squeeze()
loss = F.binary_cross_entropy_with_logits(logits, labels)
loss.backward()
optimizer.step()
return loss.detach(), labels.size(0)
def train_helper(dataloader, model, optimizer, device):
model.train() # Set the model to training mode
total_loss = torch.zeros(1, device=device) # Accumulator for the total loss
total_samples = 0 # Accumulator for the total number of samples processed
start = time.time()
for step, minibatch in tqdm(enumerate(dataloader), "Training"):
loss, num_samples = train_step(minibatch, optimizer, model)
total_loss += loss * num_samples
total_samples += num_samples
if step + 1 == args.early_stop:
break
train_loss = total_loss / total_samples
end = time.time()
return train_loss, end - start
def train(dataloader, model, device):
#####################################################################
# (HIGHLIGHT) Train the model for one epoch.
#
# - Iterates over the data loader, fetching mini-batches of graph data.
# - For each mini-batch, it performs a forward pass, computes loss, and
# updates the model parameters.
# - The function returns the average loss and accuracy for the epoch.
#
# Parameters:
# dataloader: DataLoader that provides mini-batches of graph data.
# model: The GraphSAGE model.
# device: The device (CPU/GPU) to run the training on.
#####################################################################
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
for epoch in range(args.epochs):
train_loss, duration = train_helper(
dataloader, model, optimizer, device
)
print(
f"Epoch {epoch:02d}, Loss: {train_loss.item():.4f}, "
f"Time: {duration}s"
)
def parse_args():
parser = argparse.ArgumentParser(
description="Which dataset are you going to use?"
)
parser.add_argument(
"--epochs", type=int, default=10, help="Number of training epochs."
)
parser.add_argument(
"--lr",
type=float,
default=0.003,
help="Learning rate for optimization.",
)
parser.add_argument("--neg-ratio", type=int, default=1)
parser.add_argument("--train-batch-size", type=int, default=512)
parser.add_argument("--eval-batch-size", type=int, default=1024)
parser.add_argument(
"--batch-size", type=int, default=1024, help="Batch size for training."
)
parser.add_argument(
"--num-workers",
type=int,
default=0,
help="Number of workers for data loading.",
)
parser.add_argument(
"--early-stop",
type=int,
default=0,
help="0 means no early stop, otherwise stop at the input-th step",
)
parser.add_argument(
"--dataset",
type=str,
default="ogbl-citation2",
choices=["ogbl-citation2"],
help="The dataset we can use for link prediction. Currently"
" only ogbl-citation2 dataset is supported.",
)
parser.add_argument(
"--fanout",
type=str,
default="10,10,10",
help="Fan-out of neighbor sampling. It is IMPORTANT to keep len(fanout)"
" identical with the number of layers in your model. Default: 10,10,10",
)
parser.add_argument(
"--exclude-edges",
type=bool,
default=True,
help="Whether to exclude reverse edges during sampling. Default: True",
)
parser.add_argument(
"--mode",
default="pinned-pinned-cuda",
choices=[
"cpu-cpu-cpu",
"cpu-cpu-cuda",
"cpu-pinned-cuda",
"pinned-pinned-cuda",
"cuda-pinned-cuda",
"cuda-cuda-cuda",
],
help="Graph storage - feature storage - Train device: 'cpu' for CPU and RAM,"
" 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
)
parser.add_argument(
"--gpu-cache-size",
type=int,
default=0,
help="The capacity of the GPU cache in bytes.",
)
parser.add_argument(
"--sample-mode",
default="sample_neighbor",
choices=["sample_neighbor", "sample_layer_neighbor"],
help="The sampling function when doing layerwise sampling.",
)
parser.add_argument("--precision", type=str, default="high")
return parser.parse_args()
def main():
torch.set_float32_matmul_precision(args.precision)
if not torch.cuda.is_available():
args.mode = "cpu-cpu-cpu"
print(f"Training in {args.mode} mode.")
args.graph_device, args.feature_device, args.device = args.mode.split("-")
args.overlap_feature_fetch = args.feature_device == "pinned"
args.overlap_graph_fetch = args.graph_device == "pinned"
# Load and preprocess dataset.
print("Loading data...")
dataset = gb.BuiltinDataset(args.dataset).load()
# Move the dataset to the selected storage.
graph = (
dataset.graph.pin_memory_()
if args.graph_device == "pinned"
else dataset.graph.to(args.graph_device)
)
features = (
dataset.feature.pin_memory_()
if args.feature_device == "pinned"
else dataset.feature.to(args.feature_device)
)
train_set = dataset.tasks[0].train_set
valid_set = dataset.tasks[0].validation_set
test_set = dataset.tasks[0].test_set
all_nodes_set = dataset.all_nodes_set
args.fanout = list(map(int, args.fanout.split(",")))
if args.gpu_cache_size > 0 and args.feature_device != "cuda":
features._features[("node", None, "feat")] = gb.gpu_cached_feature(
features._features[("node", None, "feat")],
args.gpu_cache_size,
)
train_dataloader = create_dataloader(
graph=graph,
features=features,
itemset=train_set,
batch_size=args.train_batch_size,
fanout=args.fanout,
device=args.device,
job="train",
)
in_channels = features.size("node", None, "feat")[0]
hidden_channels = 256
model = GraphSAGE(in_channels, hidden_channels, len(args.fanout)).to(
args.device
)
assert len(args.fanout) == len(model.layers)
train(train_dataloader, model, args.device)
# Test the model.
print("Testing...")
valid_mrr, test_mrr = evaluate(
model,
graph,
features,
all_nodes_set,
valid_set,
test_set,
)
print(
f"Validation MRR {valid_mrr.item():.4f}, Test MRR {test_mrr.item():.4f}"
)
if __name__ == "__main__":
args = parse_args()
main()
examples/graphbolt/pyg/multigpu/node_classification.py
+482
View File
@@ -0,0 +1,482 @@
"""
This script demonstrates node classification with GraphSAGE on large graphs,
merging GraphBolt (GB) and PyTorch Geometric (PyG). GraphBolt efficiently manages
data loading for large datasets, crucial for mini-batch processing. Post data
loading, PyG's user-friendly framework takes over for training, showcasing seamless
integration with GraphBolt. This combination offers an efficient alternative to
traditional Deep Graph Library (DGL) methods, highlighting adaptability and
scalability in handling large-scale graph data for diverse real-world applications.
Key Features:
- Implements the GraphSAGE model, a scalable GNN, for node classification on large graphs.
- Utilizes GraphBolt, an efficient framework for large-scale graph data processing.
- Integrates with PyTorch Geometric for building and training the GraphSAGE model.
- The script is well-documented, providing clear explanations at each step.
This flowchart describes the main functional sequence of the provided example.
main:
main
│
├───> Load and preprocess dataset (GraphBolt)
│ │
│ └───> Utilize GraphBolt's BuiltinDataset for dataset handling
│
├───> Instantiate the SAGE model (PyTorch Geometric)
│ │
│ └───> Define the GraphSAGE model architecture
│
├───> Train the model
│ │
│ ├───> Mini-Batch Processing with GraphBolt
│ │ │
│ │ └───> Efficient handling of mini-batches using GraphBolt's utilities
│ │
│ └───> Training Loop
│ │
│ ├───> Forward and backward passes
│ │
│ └───> Parameters optimization
│
└───> Evaluate the model
│
└───> Performance assessment on validation and test datasets
│
└───> Accuracy and other relevant metrics calculation
"""
import argparse
import os
import time
import dgl.graphbolt as gb
import torch
# For torch.compile until https://github.com/pytorch/pytorch/issues/121197 is
# resolved.
import torch._inductor.codecache
torch._dynamo.config.cache_size_limit = 32
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from tqdm import tqdm
def accuracy(out, labels):
assert out.ndim == 2
assert out.size(0) == labels.size(0)
assert labels.ndim == 1 or (labels.ndim == 2 and labels.size(1) == 1)
labels = labels.flatten()
predictions = torch.argmax(out, 1)
return (labels == predictions).sum(dtype=torch.float64) / labels.size(0)
class GraphSAGE(torch.nn.Module):
#####################################################################
# (HIGHLIGHT) Define the GraphSAGE model architecture.
#
# - This class inherits from `torch.nn.Module`.
# - Two convolutional layers are created using the SAGEConv class from PyG.
# - 'in_size', 'hidden_size', 'out_size' are the sizes of
# the input, hidden, and output features, respectively.
# - The forward method defines the computation performed at every call.
#####################################################################
def __init__(self, in_size, hidden_size, out_size, n_layers, cooperative):
super(GraphSAGE, self).__init__()
self.layers = torch.nn.ModuleList()
sizes = [in_size] + [hidden_size] * (n_layers - 1) + [out_size]
for i in range(n_layers):
self.layers.append(SAGEConv(sizes[i], sizes[i + 1]))
self.hidden_size = hidden_size
self.out_size = out_size
self.cooperative = cooperative
def forward(self, minibatch, x):
subgraphs = minibatch.sampled_subgraphs
h = x
for i, (layer, subgraph) in enumerate(zip(self.layers, subgraphs)):
#####################################################################
# (HIGHLIGHT) Convert given features to be consumed by a PyG layer.
#
# PyG layers have two modes, bipartite and normal. We slice the
# given features to get src and dst features to use the PyG layers
# in the more efficient bipartite mode.
#####################################################################
if i != 0 and self.cooperative:
h = gb.CooperativeConvFunction.apply(subgraph, h)
h, edge_index, size = subgraph.to_pyg(h)
h = layer(h, edge_index, size=size)
if i != len(subgraphs) - 1:
h = F.relu(h)
if self.cooperative:
h = gb.CooperativeConvFunction.apply(minibatch, h)
h = h[minibatch.compacted_seeds]
return h
def create_dataloader(
args, graph, features, itemset, batch_size, fanout, device, job
):
#####################################################################
# (HIGHLIGHT) Create a data loader for efficiently loading graph data.
#
# - 'ItemSampler' samples mini-batches of node IDs from the dataset.
# - 'CopyTo' copies the fetched data to the specified device.
# - 'sample_neighbor' performs neighbor sampling on the graph.
# - 'FeatureFetcher' fetches node features based on the sampled subgraph.
#####################################################################
# Create a datapipe for mini-batch sampling with a specific neighbor fanout.
# Here, [10, 10, 10] specifies the number of neighbors sampled for each node at each layer.
# We're using `sample_neighbor` for consistency with DGL's sampling API.
# Note: GraphBolt offers additional sampling methods, such as `sample_layer_neighbor`,
# which could provide further optimization and efficiency for GNN training.
# Users are encouraged to explore these advanced features for potentially improved performance.
# Initialize an ItemSampler to sample mini-batches from the dataset.
datapipe = gb.DistributedItemSampler(
itemset,
batch_size=batch_size,
shuffle=(job == "train"),
drop_last=(job == "train"),
drop_uneven_inputs=True,
)
need_copy = True
# Copy the data to the specified device.
if args.graph_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
# Sample neighbors for each node in the mini-batch.
datapipe = getattr(datapipe, args.sample_mode)(
graph,
fanout if job != "infer" else [-1],
overlap_fetch=args.overlap_graph_fetch,
num_gpu_cached_edges=args.num_gpu_cached_edges,
gpu_cache_threshold=args.gpu_graph_caching_threshold,
cooperative=args.cooperative,
asynchronous=args.graph_device != "cpu",
)
# Copy the data to the specified device.
if args.feature_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
# Fetch node features for the sampled subgraph.
datapipe = datapipe.fetch_feature(
features,
node_feature_keys=["feat"],
overlap_fetch=args.overlap_feature_fetch,
cooperative=args.cooperative,
)
# Copy the data to the specified device.
if need_copy:
datapipe = datapipe.copy_to(device=device)
# Create and return a DataLoader to handle data loading.
return gb.DataLoader(datapipe, num_workers=args.num_workers)
def weighted_reduce(tensor, weight, dst=0):
########################################################################
# (HIGHLIGHT) Collect accuracy and loss values from sub-processes and
# obtain overall average values.
#
# `torch.distributed.reduce` is used to reduce tensors from all the
# sub-processes to a specified process, ReduceOp.SUM is used by default.
#
# Because the GPUs may have differing numbers of processed items, we
# perform a weighted mean to calculate the exact loss and accuracy.
########################################################################
dist.reduce(tensor=tensor, dst=dst)
weight = torch.tensor(weight, device=tensor.device)
dist.reduce(tensor=weight, dst=dst)
return tensor / weight
@torch.compile
def train_step(minibatch, optimizer, model, loss_fn):
node_features = minibatch.node_features["feat"]
labels = minibatch.labels
optimizer.zero_grad()
out = model(minibatch, node_features)
loss = loss_fn(out, labels)
num_correct = accuracy(out, labels) * labels.size(0)
loss.backward()
optimizer.step()
return loss.detach(), num_correct, labels.size(0)
def train_helper(rank, dataloader, model, optimizer, loss_fn, device):
model.train() # Set the model to training mode
total_loss = torch.zeros(1, device=device) # Accumulator for the total loss
# Accumulator for the total number of correct predictions
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0 # Accumulator for the total number of samples processed
num_batches = 0 # Counter for the number of mini-batches processed
start = time.time()
for minibatch in tqdm(dataloader, "Training") if rank == 0 else dataloader:
loss, num_correct, num_samples = train_step(
minibatch, optimizer, model, loss_fn
)
total_loss += loss
total_correct += num_correct
total_samples += num_samples
num_batches += 1
train_loss = weighted_reduce(total_loss, num_batches)
train_acc = weighted_reduce(total_correct, total_samples)
end = time.time()
return train_loss, train_acc, end - start
def train(args, rank, train_dataloader, valid_dataloader, model, device):
#####################################################################
# (HIGHLIGHT) Train the model for one epoch.
#
# - Iterates over the data loader, fetching mini-batches of graph data.
# - For each mini-batch, it performs a forward pass, computes loss, and
# updates the model parameters.
# - The function returns the average loss and accuracy for the epoch.
#
# Parameters:
# model: The GraphSAGE model.
# dataloader: DataLoader that provides mini-batches of graph data.
# optimizer: Optimizer used for updating model parameters.
# loss_fn: Loss function used for training.
# device: The device (CPU/GPU) to run the training on.
#####################################################################
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
loss_fn = torch.nn.CrossEntropyLoss()
for epoch in range(args.epochs):
train_loss, train_acc, duration = train_helper(
rank,
train_dataloader,
model,
optimizer,
loss_fn,
device,
)
val_acc = evaluate(rank, model, valid_dataloader, device)
if rank == 0:
print(
f"Epoch {epoch:02d}, Loss: {train_loss.item():.4f}, "
f"Approx. Train: {train_acc.item():.4f}, "
f"Approx. Val: {val_acc.item():.4f}, "
f"Time: {duration}s"
)
@torch.compile
def evaluate_step(minibatch, model):
node_features = minibatch.node_features["feat"]
labels = minibatch.labels
out = model(minibatch, node_features)
num_correct = accuracy(out, labels) * labels.size(0)
return num_correct, labels.size(0)
@torch.no_grad()
def evaluate(rank, model, dataloader, device):
model.eval()
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0
for minibatch in (
tqdm(dataloader, "Evaluating") if rank == 0 else dataloader
):
num_correct, num_samples = evaluate_step(minibatch, model)
total_correct += num_correct
total_samples += num_samples
return weighted_reduce(total_correct, total_samples)
def parse_args():
parser = argparse.ArgumentParser(
description="Which dataset are you going to use?"
)
parser.add_argument(
"--epochs", type=int, default=10, help="Number of training epochs."
)
parser.add_argument(
"--lr",
type=float,
default=0.003,
help="Learning rate for optimization.",
)
parser.add_argument(
"--batch-size", type=int, default=1024, help="Batch size for training."
)
parser.add_argument(
"--num-workers",
type=int,
default=0,
help="Number of workers for data loading.",
)
parser.add_argument(
"--dataset",
type=str,
default="ogbn-products",
choices=[
"ogbn-arxiv",
"ogbn-products",
"ogbn-papers100M",
"igb-hom-tiny",
"igb-hom-small",
"igb-hom-medium",
"igb-hom-large",
"igb-hom",
],
help="The dataset we can use for node classification example. Currently"
" ogbn-products, ogbn-arxiv, ogbn-papers100M and"
" igb-hom-[tiny|small|medium|large] and igb-hom datasets are supported.",
)
parser.add_argument(
"--fanout",
type=str,
default="10,10,10",
help="Fan-out of neighbor sampling. It is IMPORTANT to keep len(fanout)"
" identical with the number of layers in your model. Default: 10,10,10",
)
parser.add_argument(
"--mode",
default="pinned-pinned-cuda",
choices=[
"pinned-pinned-cuda",
"cuda-pinned-cuda",
"cuda-cuda-cuda",
],
help="Graph storage - feature storage - Train device: 'cpu' for CPU and RAM,"
" 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
)
parser.add_argument(
"--gpu-cache-size",
type=int,
default=0,
help="The capacity of the GPU cache in bytes.",
)
parser.add_argument(
"--sample-mode",
default="sample_neighbor",
choices=["sample_neighbor", "sample_layer_neighbor"],
help="The sampling function when doing layerwise sampling.",
)
parser.add_argument(
"--num-gpu-cached-edges",
type=int,
default=0,
help="The number of edges to be cached from the graph on the GPU.",
)
parser.add_argument(
"--gpu-graph-caching-threshold",
type=int,
default=1,
help="The number of accesses after which a vertex neighborhood will be cached.",
)
parser.add_argument("--precision", type=str, default="medium")
parser.add_argument(
"--cooperative",
action="store_true",
help="Enables Cooperative Minibatching from arXiv:2310.12403.",
)
return parser.parse_args()
def run(rank, world_size, args, dataset):
# Set up multiprocessing environment.
torch.cuda.set_device(rank)
dist.init_process_group(
init_method="tcp://127.0.0.1:12345",
rank=rank,
world_size=world_size,
)
print(f"Training in {args.mode} mode.")
args.graph_device, args.feature_device, args.device = args.mode.split("-")
args.overlap_feature_fetch = args.feature_device == "pinned"
args.overlap_graph_fetch = args.graph_device == "pinned"
# Move the dataset to the selected storage.
graph = (
dataset.graph.pin_memory_()
if args.graph_device == "pinned"
else dataset.graph.to(args.graph_device)
)
features = (
dataset.feature.pin_memory_()
if args.feature_device == "pinned"
else dataset.feature.to(args.feature_device)
)
train_set = dataset.tasks[0].train_set
valid_set = dataset.tasks[0].validation_set
args.fanout = list(map(int, args.fanout.split(",")))
num_classes = dataset.tasks[0].metadata["num_classes"]
if args.gpu_cache_size > 0 and args.feature_device != "cuda":
features._features[("node", None, "feat")] = gb.gpu_cached_feature(
features._features[("node", None, "feat")],
args.gpu_cache_size,
)
train_dataloader, valid_dataloader = (
create_dataloader(
args,
graph=graph,
features=features,
itemset=itemset,
batch_size=args.batch_size,
fanout=args.fanout,
device=args.device,
job=job,
)
for itemset, job in zip([train_set, valid_set], ["train", "evaluate"])
)
in_channels = features.size("node", None, "feat")[0]
hidden_channels = 256
model = GraphSAGE(
in_channels,
hidden_channels,
num_classes,
len(args.fanout),
args.cooperative,
).to(args.device)
assert len(args.fanout) == len(model.layers)
model = torch.nn.parallel.DistributedDataParallel(model)
train(args, rank, train_dataloader, valid_dataloader, model, args.device)
dist.destroy_process_group()
if __name__ == "__main__":
args = parse_args()
if not torch.cuda.is_available():
print("Multi-GPU training requires GPUs.")
exit(0)
torch.set_float32_matmul_precision(args.precision)
# Load and preprocess dataset.
print("Loading data...")
dataset = gb.BuiltinDataset(args.dataset).load()
world_size = torch.cuda.device_count()
# Thread limiting to avoid resource competition.
os.environ["OMP_NUM_THREADS"] = str(mp.cpu_count() // 2 // world_size)
mp.set_sharing_strategy("file_system")
mp.spawn(
run,
args=(world_size, args, dataset),
nprocs=world_size,
join=True,
)
examples/graphbolt/pyg/node_classification.py
+280
View File
@@ -0,0 +1,280 @@
"""
This script demonstrates node classification with GraphSAGE on large graphs,
merging GraphBolt (GB) and PyTorch Geometric (PyG). GraphBolt efficiently
manages data loading for large datasets, crucial for mini-batch processing.
Post data loading, PyG's user-friendly framework takes over for training,
showcasing seamless integration with GraphBolt. This combination offers an
efficient alternative to traditional Deep Graph Library (DGL) methods,
highlighting adaptability and scalability in handling large-scale graph data
for diverse real-world applications.
Key Features:
- Implements the GraphSAGE model, a scalable GNN, for node classification on
large graphs.
- Utilizes GraphBolt, an efficient framework for large-scale graph data processing.
- Integrates with PyTorch Geometric for building and training the GraphSAGE model.
- The script is well-documented, providing clear explanations at each step.
This flowchart describes the main functional sequence of the provided example.
main:
main
│
├───> Load and preprocess dataset (GraphBolt)
│ │
│ └───> Utilize GraphBolt's BuiltinDataset for dataset handling
│
├───> Instantiate the SAGE model (PyTorch Geometric)
│ │
│ └───> Define the GraphSAGE model architecture
│
├───> Train the model
│ │
│ ├───> Mini-Batch Processing with GraphBolt
│ │ │
│ │ └───> Efficient handling of mini-batches using GraphBolt's utilities
│ │
│ └───> Training Loop
│ │
│ ├───> Forward and backward passes
│ │
│ ├───> Convert GraphBolt MiniBatch to PyG Data
│ │
│ └───> Parameters optimization
│
└───> Evaluate the model
│
└───> Performance assessment on validation and test datasets
│
└───> Accuracy and other relevant metrics calculation
"""
import argparse
import dgl.graphbolt as gb
import torch
import torch.nn.functional as F
import torchmetrics.functional as MF
from torch_geometric.nn import SAGEConv
from tqdm import tqdm
class GraphSAGE(torch.nn.Module):
#####################################################################
# (HIGHLIGHT) Define the GraphSAGE model architecture.
#
# - This class inherits from `torch.nn.Module`.
# - Two convolutional layers are created using the SAGEConv class from PyG.
# - 'in_size', 'hidden_size', 'out_size' are the sizes of
# the input, hidden, and output features, respectively.
# - The forward method defines the computation performed at every call.
# - It's adopted from the official PyG example which can be found at
# https://github.com/pyg-team/pytorch_geometric/blob/master/examples/ogbn_products_sage.py
#####################################################################
def __init__(self, in_size, hidden_size, out_size):
super(GraphSAGE, self).__init__()
self.layers = torch.nn.ModuleList()
self.layers.append(SAGEConv(in_size, hidden_size))
self.layers.append(SAGEConv(hidden_size, hidden_size))
self.layers.append(SAGEConv(hidden_size, out_size))
def forward(self, x, edge_index):
for i, layer in enumerate(self.layers):
x = layer(x, edge_index)
if i != len(self.layers) - 1:
x = x.relu()
x = F.dropout(x, p=0.5, training=self.training)
return x
def inference(self, dataloader, x_all, device):
"""Conduct layer-wise inference to get all the node embeddings."""
for i, layer in tqdm(enumerate(self.layers), "inference"):
xs = []
for minibatch in dataloader:
# Call `to_pyg_data` to convert GB Minibatch to PyG Data.
pyg_data = minibatch.to_pyg_data()
n_id = pyg_data.n_id.to("cpu")
x = x_all[n_id].to(device)
edge_index = pyg_data.edge_index
x = layer(x, edge_index)
x = x[: pyg_data.batch_size]
if i != len(self.layers) - 1:
x = x.relu()
xs.append(x.cpu())
x_all = torch.cat(xs, dim=0)
return x_all
def create_dataloader(
dataset_set, graph, feature, batch_size, fanout, device, job
):
# Initialize an ItemSampler to sample mini-batches from the dataset.
datapipe = gb.ItemSampler(
dataset_set,
batch_size=batch_size,
shuffle=(job == "train"),
drop_last=(job == "train"),
)
# Sample neighbors for each node in the mini-batch.
datapipe = datapipe.sample_neighbor(
graph, fanout if job != "infer" else [-1]
)
# Copy the data to the specified device.
datapipe = datapipe.copy_to(device=device)
# Fetch node features for the sampled subgraph.
datapipe = datapipe.fetch_feature(feature, node_feature_keys=["feat"])
# Create and return a DataLoader to handle data loading.
dataloader = gb.DataLoader(datapipe, num_workers=0)
return dataloader
def train(model, dataloader, optimizer):
model.train() # Set the model to training mode
total_loss = 0 # Accumulator for the total loss
total_correct = 0 # Accumulator for the total number of correct predictions
total_samples = 0 # Accumulator for the total number of samples processed
num_batches = 0 # Counter for the number of mini-batches processed
for _, minibatch in tqdm(enumerate(dataloader), "training"):
#####################################################################
# (HIGHLIGHT) Convert GraphBolt MiniBatch to PyG Data class.
#
# Call `MiniBatch.to_pyg_data()` and it will return a PyG Data class
# with necessary data and information.
#####################################################################
pyg_data = minibatch.to_pyg_data()
optimizer.zero_grad()
out = model(pyg_data.x, pyg_data.edge_index)[: pyg_data.y.shape[0]]
y = pyg_data.y
loss = F.cross_entropy(out, y)
loss.backward()
optimizer.step()
total_loss += float(loss)
total_correct += int(out.argmax(dim=-1).eq(y).sum())
total_samples += y.shape[0]
num_batches += 1
avg_loss = total_loss / num_batches
avg_accuracy = total_correct / total_samples
return avg_loss, avg_accuracy
@torch.no_grad()
def evaluate(model, dataloader, num_classes):
model.eval()
y_hats = []
ys = []
for _, minibatch in tqdm(enumerate(dataloader), "evaluating"):
pyg_data = minibatch.to_pyg_data()
out = model(pyg_data.x, pyg_data.edge_index)[: pyg_data.y.shape[0]]
y = pyg_data.y
y_hats.append(out)
ys.append(y)
return MF.accuracy(
torch.cat(y_hats),
torch.cat(ys),
task="multiclass",
num_classes=num_classes,
)
@torch.no_grad()
def layerwise_infer(
model, infer_dataloader, test_set, feature, num_classes, device
):
model.eval()
features = feature.read("node", None, "feat")
pred = model.inference(infer_dataloader, features, device)
pred = pred[test_set._items[0]]
label = test_set._items[1].to(pred.device)
return MF.accuracy(
pred,
label,
task="multiclass",
num_classes=num_classes,
)
def main():
parser = argparse.ArgumentParser(
description="Which dataset are you going to use?"
)
parser.add_argument(
"--dataset",
type=str,
default="ogbn-products",
help='Name of the dataset to use (e.g., "ogbn-products", "ogbn-arxiv")',
)
parser.add_argument(
"--epochs", type=int, default=10, help="Number of training epochs."
)
parser.add_argument(
"--batch-size", type=int, default=1024, help="Batch size for training."
)
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dataset_name = args.dataset
dataset = gb.BuiltinDataset(dataset_name).load()
graph = dataset.graph
feature = dataset.feature.pin_memory_()
train_set = dataset.tasks[0].train_set
valid_set = dataset.tasks[0].validation_set
test_set = dataset.tasks[0].test_set
all_nodes_set = dataset.all_nodes_set
num_classes = dataset.tasks[0].metadata["num_classes"]
train_dataloader = create_dataloader(
train_set,
graph,
feature,
args.batch_size,
[5, 10, 15],
device,
job="train",
)
valid_dataloader = create_dataloader(
valid_set,
graph,
feature,
args.batch_size,
[5, 10, 15],
device,
job="evaluate",
)
infer_dataloader = create_dataloader(
all_nodes_set,
graph,
feature,
4 * args.batch_size,
[-1],
device,
job="infer",
)
in_channels = feature.size("node", None, "feat")[0]
hidden_channels = 256
model = GraphSAGE(in_channels, hidden_channels, num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.003)
for epoch in range(args.epochs):
train_loss, train_accuracy = train(model, train_dataloader, optimizer)
valid_accuracy = evaluate(model, valid_dataloader, num_classes)
print(
f"Epoch {epoch}, Train Loss: {train_loss:.4f}, "
f"Train Accuracy: {train_accuracy:.4f}, "
f"Valid Accuracy: {valid_accuracy:.4f}"
)
test_accuracy = layerwise_infer(
model, infer_dataloader, test_set, feature, num_classes, device
)
print(f"Test Accuracy: {test_accuracy:.4f}")
if __name__ == "__main__":
main()
examples/graphbolt/pyg/node_classification_advanced.py
+477
View File
@@ -0,0 +1,477 @@
"""
This script demonstrates node classification with GraphSAGE on large graphs,
merging GraphBolt (GB) and PyTorch Geometric (PyG). GraphBolt efficiently manages
data loading for large datasets, crucial for mini-batch processing. Post data
loading, PyG's user-friendly framework takes over for training, showcasing seamless
integration with GraphBolt. This combination offers an efficient alternative to
traditional Deep Graph Library (DGL) methods, highlighting adaptability and
scalability in handling large-scale graph data for diverse real-world applications.
Key Features:
- Implements the GraphSAGE model, a scalable GNN, for node classification on large graphs.
- Utilizes GraphBolt, an efficient framework for large-scale graph data processing.
- Integrates with PyTorch Geometric for building and training the GraphSAGE model.
- The script is well-documented, providing clear explanations at each step.
This flowchart describes the main functional sequence of the provided example.
main:
main
│
├───> Load and preprocess dataset (GraphBolt)
│ │
│ └───> Utilize GraphBolt's BuiltinDataset for dataset handling
│
├───> Instantiate the SAGE model (PyTorch Geometric)
│ │
│ └───> Define the GraphSAGE model architecture
│
├───> Train the model
│ │
│ ├───> Mini-Batch Processing with GraphBolt
│ │ │
│ │ └───> Efficient handling of mini-batches using GraphBolt's utilities
│ │
│ └───> Training Loop
│ │
│ ├───> Forward and backward passes
│ │
│ └───> Parameters optimization
│
└───> Evaluate the model
│
└───> Performance assessment on validation and test datasets
│
└───> Accuracy and other relevant metrics calculation
"""
import argparse
import time
import dgl.graphbolt as gb
import torch
# For torch.compile until https://github.com/pytorch/pytorch/issues/121197 is
# resolved.
import torch._inductor.codecache
torch._dynamo.config.cache_size_limit = 32
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from tqdm import tqdm
def accuracy(out, labels):
assert out.ndim == 2
assert out.size(0) == labels.size(0)
assert labels.ndim == 1 or (labels.ndim == 2 and labels.size(1) == 1)
labels = labels.flatten()
predictions = torch.argmax(out, 1)
return (labels == predictions).sum(dtype=torch.float64) / labels.size(0)
class GraphSAGE(torch.nn.Module):
#####################################################################
# (HIGHLIGHT) Define the GraphSAGE model architecture.
#
# - This class inherits from `torch.nn.Module`.
# - Two convolutional layers are created using the SAGEConv class from PyG.
# - 'in_size', 'hidden_size', 'out_size' are the sizes of
# the input, hidden, and output features, respectively.
# - The forward method defines the computation performed at every call.
#####################################################################
def __init__(self, in_size, hidden_size, out_size, n_layers):
super(GraphSAGE, self).__init__()
self.layers = torch.nn.ModuleList()
sizes = [in_size] + [hidden_size] * (n_layers - 1) + [out_size]
for i in range(n_layers):
self.layers.append(SAGEConv(sizes[i], sizes[i + 1]))
self.hidden_size = hidden_size
self.out_size = out_size
def forward(self, subgraphs, x):
h = x
for i, (layer, subgraph) in enumerate(zip(self.layers, subgraphs)):
#####################################################################
# (HIGHLIGHT) Convert given features to be consumed by a PyG layer.
#
# PyG layers have two modes, bipartite and normal. We slice the
# given features to get src and dst features to use the PyG layers
# in the more efficient bipartite mode.
#####################################################################
h, edge_index, size = subgraph.to_pyg(h)
h = layer(h, edge_index, size=size)
if i != len(subgraphs) - 1:
h = F.relu(h)
return h
def inference(self, graph, features, dataloader, storage_device):
"""Conduct layer-wise inference to get all the node embeddings."""
pin_memory = storage_device == "pinned"
buffer_device = torch.device("cpu" if pin_memory else storage_device)
for layer_idx, layer in enumerate(self.layers):
is_last_layer = layer_idx == len(self.layers) - 1
y = torch.empty(
graph.total_num_nodes,
self.out_size if is_last_layer else self.hidden_size,
dtype=torch.float32,
device=buffer_device,
pin_memory=pin_memory,
)
for data in tqdm(dataloader, "Inferencing"):
# len(data.sampled_subgraphs) = 1
h, edge_index, size = data.sampled_subgraphs[0].to_pyg(
data.node_features["feat"]
)
hidden_x = layer(h, edge_index, size=size)
if not is_last_layer:
hidden_x = F.relu(hidden_x)
# By design, our output nodes are contiguous.
y[data.seeds[0] : data.seeds[-1] + 1] = hidden_x.to(
buffer_device
)
if not is_last_layer:
features.update("node", None, "feat", y)
return y
def create_dataloader(
graph, features, itemset, batch_size, fanout, device, job
):
#####################################################################
# (HIGHLIGHT) Create a data loader for efficiently loading graph data.
#
# - 'ItemSampler' samples mini-batches of node IDs from the dataset.
# - 'CopyTo' copies the fetched data to the specified device.
# - 'sample_neighbor' performs neighbor sampling on the graph.
# - 'FeatureFetcher' fetches node features based on the sampled subgraph.
#####################################################################
# Create a datapipe for mini-batch sampling with a specific neighbor fanout.
# Here, [10, 10, 10] specifies the number of neighbors sampled for each node at each layer.
# We're using `sample_neighbor` for consistency with DGL's sampling API.
# Note: GraphBolt offers additional sampling methods, such as `sample_layer_neighbor`,
# which could provide further optimization and efficiency for GNN training.
# Users are encouraged to explore these advanced features for potentially improved performance.
# Initialize an ItemSampler to sample mini-batches from the dataset.
datapipe = gb.ItemSampler(
itemset,
batch_size=batch_size,
shuffle=(job == "train"),
drop_last=(job == "train"),
)
need_copy = True
# Copy the data to the specified device.
if args.graph_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
# Sample neighbors for each node in the mini-batch.
datapipe = getattr(datapipe, args.sample_mode)(
graph,
fanout if job != "infer" else [-1],
overlap_fetch=args.overlap_graph_fetch,
num_gpu_cached_edges=args.num_gpu_cached_edges,
gpu_cache_threshold=args.gpu_graph_caching_threshold,
asynchronous=args.graph_device != "cpu",
)
# Copy the data to the specified device.
if args.feature_device != "cpu" and need_copy:
datapipe = datapipe.copy_to(device=device)
need_copy = False
# Fetch node features for the sampled subgraph.
datapipe = datapipe.fetch_feature(
features,
node_feature_keys=["feat"],
overlap_fetch=args.overlap_feature_fetch,
)
# Copy the data to the specified device.
if need_copy:
datapipe = datapipe.copy_to(device=device)
# Create and return a DataLoader to handle data loading.
return gb.DataLoader(datapipe, num_workers=args.num_workers)
@torch.compile
def train_step(minibatch, optimizer, model, loss_fn):
node_features = minibatch.node_features["feat"]
labels = minibatch.labels
optimizer.zero_grad()
out = model(minibatch.sampled_subgraphs, node_features)
loss = loss_fn(out, labels)
num_correct = accuracy(out, labels) * labels.size(0)
loss.backward()
optimizer.step()
return loss.detach(), num_correct, labels.size(0)
def train_helper(dataloader, model, optimizer, loss_fn, device):
model.train() # Set the model to training mode
total_loss = torch.zeros(1, device=device) # Accumulator for the total loss
# Accumulator for the total number of correct predictions
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0 # Accumulator for the total number of samples processed
num_batches = 0 # Counter for the number of mini-batches processed
start = time.time()
for minibatch in tqdm(dataloader, "Training"):
loss, num_correct, num_samples = train_step(
minibatch, optimizer, model, loss_fn
)
total_loss += loss
total_correct += num_correct
total_samples += num_samples
num_batches += 1
train_loss = total_loss / num_batches
train_acc = total_correct / total_samples
end = time.time()
return train_loss, train_acc, end - start
def train(train_dataloader, valid_dataloader, model, device):
#####################################################################
# (HIGHLIGHT) Train the model for one epoch.
#
# - Iterates over the data loader, fetching mini-batches of graph data.
# - For each mini-batch, it performs a forward pass, computes loss, and
# updates the model parameters.
# - The function returns the average loss and accuracy for the epoch.
#
# Parameters:
# model: The GraphSAGE model.
# dataloader: DataLoader that provides mini-batches of graph data.
# optimizer: Optimizer used for updating model parameters.
# loss_fn: Loss function used for training.
# device: The device (CPU/GPU) to run the training on.
#####################################################################
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
loss_fn = torch.nn.CrossEntropyLoss()
for epoch in range(args.epochs):
train_loss, train_acc, duration = train_helper(
train_dataloader, model, optimizer, loss_fn, device
)
val_acc = evaluate(model, valid_dataloader, device)
print(
f"Epoch {epoch:02d}, Loss: {train_loss.item():.4f}, "
f"Approx. Train: {train_acc.item():.4f}, "
f"Approx. Val: {val_acc.item():.4f}, "
f"Time: {duration}s"
)
@torch.no_grad()
def layerwise_infer(args, graph, features, test_set, all_nodes_set, model):
model.eval()
dataloader = create_dataloader(
graph=graph,
features=features,
itemset=all_nodes_set,
batch_size=4 * args.batch_size,
fanout=[-1],
device=args.device,
job="infer",
)
pred = model.inference(graph, features, dataloader, args.feature_device)
pred = pred[test_set._items[0]]
label = test_set._items[1].to(pred.device)
return accuracy(pred, label)
@torch.compile
def evaluate_step(minibatch, model):
node_features = minibatch.node_features["feat"]
labels = minibatch.labels
out = model(minibatch.sampled_subgraphs, node_features)
num_correct = accuracy(out, labels) * labels.size(0)
return num_correct, labels.size(0)
@torch.no_grad()
def evaluate(model, dataloader, device):
model.eval()
total_correct = torch.zeros(1, dtype=torch.float64, device=device)
total_samples = 0
for minibatch in tqdm(dataloader, "Evaluating"):
num_correct, num_samples = evaluate_step(minibatch, model)
total_correct += num_correct
total_samples += num_samples
return total_correct / total_samples
def parse_args():
parser = argparse.ArgumentParser(
description="Which dataset are you going to use?"
)
parser.add_argument(
"--epochs", type=int, default=10, help="Number of training epochs."
)
parser.add_argument(
"--lr",
type=float,
default=0.003,
help="Learning rate for optimization.",
)
parser.add_argument(
"--batch-size", type=int, default=1024, help="Batch size for training."
)
parser.add_argument(
"--num-workers",
type=int,
default=0,
help="Number of workers for data loading.",
)
parser.add_argument(
"--dataset",
type=str,
default="ogbn-products",
choices=[
"ogbn-arxiv",
"ogbn-products",
"ogbn-papers100M",
"igb-hom-tiny",
"igb-hom-small",
"igb-hom-medium",
"igb-hom-large",
"igb-hom",
],
help="The dataset we can use for node classification example. Currently"
" ogbn-products, ogbn-arxiv, ogbn-papers100M and"
" igb-hom-[tiny|small|medium|large] and igb-hom datasets are supported.",
)
parser.add_argument(
"--fanout",
type=str,
default="10,10,10",
help="Fan-out of neighbor sampling. It is IMPORTANT to keep len(fanout)"
" identical with the number of layers in your model. Default: 10,10,10",
)
parser.add_argument(
"--mode",
default="pinned-pinned-cuda",
choices=[
"cpu-cpu-cpu",
"cpu-cpu-cuda",
"cpu-pinned-cuda",
"pinned-pinned-cuda",
"cuda-pinned-cuda",
"cuda-cuda-cuda",
],
help="Graph storage - feature storage - Train device: 'cpu' for CPU and RAM,"
" 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
)
parser.add_argument(
"--gpu-cache-size",
type=int,
default=0,
help="The capacity of the GPU cache in bytes.",
)
parser.add_argument(
"--sample-mode",
default="sample_neighbor",
choices=["sample_neighbor", "sample_layer_neighbor"],
help="The sampling function when doing layerwise sampling.",
)
parser.add_argument(
"--num-gpu-cached-edges",
type=int,
default=0,
help="The number of edges to be cached from the graph on the GPU.",
)
parser.add_argument(
"--gpu-graph-caching-threshold",
type=int,
default=1,
help="The number of accesses after which a vertex neighborhood will be cached.",
)
parser.add_argument("--precision", type=str, default="high")
return parser.parse_args()
def main():
torch.set_float32_matmul_precision(args.precision)
if not torch.cuda.is_available():
args.mode = "cpu-cpu-cpu"
print(f"Training in {args.mode} mode.")
args.graph_device, args.feature_device, args.device = args.mode.split("-")
args.overlap_feature_fetch = args.feature_device == "pinned"
args.overlap_graph_fetch = args.graph_device == "pinned"
# Load and preprocess dataset.
print("Loading data...")
dataset = gb.BuiltinDataset(args.dataset).load()
# Move the dataset to the selected storage.
graph = (
dataset.graph.pin_memory_()
if args.graph_device == "pinned"
else dataset.graph.to(args.graph_device)
)
features = (
dataset.feature.pin_memory_()
if args.feature_device == "pinned"
else dataset.feature.to(args.feature_device)
)
train_set = dataset.tasks[0].train_set
valid_set = dataset.tasks[0].validation_set
test_set = dataset.tasks[0].test_set
all_nodes_set = dataset.all_nodes_set
args.fanout = list(map(int, args.fanout.split(",")))
num_classes = dataset.tasks[0].metadata["num_classes"]
if args.gpu_cache_size > 0 and args.feature_device != "cuda":
features._features[("node", None, "feat")] = gb.gpu_cached_feature(
features._features[("node", None, "feat")],
args.gpu_cache_size,
)
train_dataloader, valid_dataloader = (
create_dataloader(
graph=graph,
features=features,
itemset=itemset,
batch_size=args.batch_size,
fanout=args.fanout,
device=args.device,
job=job,
)
for itemset, job in zip([train_set, valid_set], ["train", "evaluate"])
)
in_channels = features.size("node", None, "feat")[0]
hidden_channels = 256
model = GraphSAGE(
in_channels, hidden_channels, num_classes, len(args.fanout)
).to(args.device)
assert len(args.fanout) == len(model.layers)
train(train_dataloader, valid_dataloader, model, args.device)
# Test the model.
print("Testing...")
test_acc = layerwise_infer(
args,
graph,
features,
test_set,
all_nodes_set,
model,
)
print(f"Test accuracy {test_acc.item():.4f}")
if __name__ == "__main__":
args = parse_args()
main()
examples/graphbolt/quickstart/README.md
+11
View File
@@ -0,0 +1,11 @@
# Graphbolt Quickstart Tutorial
Graphbolt provides all you need to create a dataloader to train a Graph Neural Networks.
## Examples
- The [node_classification.py](https://github.com/dmlc/dgl/blob/master/examples/graphbolt/quickstart/node_classification.py)
shows how to create a Graphbolt dataloader to train a 2 layer Graph Convolutional Networks node
classification model.
- The [link_prediction.py](https://github.com/dmlc/dgl/blob/master/examples/graphbolt/quickstart/link_prediction.py)
shows how to create a Graphbolt dataloader to train a 2 layer GraphSage link prediction model.
examples/graphbolt/quickstart/link_prediction.py
+178
View File
@@ -0,0 +1,178 @@
"""
This example shows how to create a GraphBolt dataloader to sample and train a
link prediction model with the Cora dataset.
Disclaimer: Please note that the test edges are not excluded from the original
graph in the dataset, which could lead to data leakage. We are ignoring this
issue for this example because we are focused on demonstrating usability.
"""
import dgl.graphbolt as gb
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import SAGEConv
from torcheval.metrics import BinaryAUROC
############################################################################
# (HIGHLIGHT) Create a single process dataloader with dgl graphbolt package.
############################################################################
def create_dataloader(dataset, device, is_train=True):
# The second of two tasks in the dataset is link prediction.
task = dataset.tasks[1]
itemset = task.train_set if is_train else task.test_set
# Sample seed edges from the itemset.
datapipe = gb.ItemSampler(itemset, batch_size=256)
# Copy the mini-batch to the designated device for sampling and training.
datapipe = datapipe.copy_to(device)
if is_train:
# Sample negative edges for the seed edges.
datapipe = datapipe.sample_uniform_negative(
dataset.graph, negative_ratio=1
)
# Sample neighbors for the seed nodes.
datapipe = datapipe.sample_neighbor(dataset.graph, fanouts=[4, 2])
# Exclude seed edges from the subgraph.
datapipe = datapipe.transform(gb.exclude_seed_edges)
else:
# Sample neighbors for the seed nodes.
datapipe = datapipe.sample_neighbor(dataset.graph, fanouts=[-1, -1])
# Fetch features for sampled nodes.
datapipe = datapipe.fetch_feature(
dataset.feature, node_feature_keys=["feat"]
)
# Initiate the dataloader for the datapipe.
return gb.DataLoader(datapipe)
class GraphSAGE(nn.Module):
def __init__(self, in_size, hidden_size=16):
super().__init__()
self.layers = nn.ModuleList()
self.layers.append(SAGEConv(in_size, hidden_size, "mean"))
self.layers.append(SAGEConv(hidden_size, hidden_size, "mean"))
self.predictor = nn.Sequential(
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, 1),
)
def forward(self, blocks, x):
hidden_x = x
for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)):
hidden_x = layer(block, hidden_x)
is_last_layer = layer_idx == len(self.layers) - 1
if not is_last_layer:
hidden_x = F.relu(hidden_x)
return hidden_x
@torch.no_grad()
def evaluate(model, dataset, device):
model.eval()
dataloader = create_dataloader(dataset, device, is_train=False)
logits = []
labels = []
for step, data in enumerate(dataloader):
# Get node pairs with labels for loss calculation.
compacted_seeds = data.compacted_seeds.T
label = data.labels
# The features of sampled nodes.
x = data.node_features["feat"]
# Forward.
y = model(data.blocks, x)
logit = (
model.predictor(
y[compacted_seeds[0].long()] * y[compacted_seeds[1].long()]
)
.squeeze()
.detach()
)
logits.append(logit)
labels.append(label)
logits = torch.cat(logits, dim=0)
labels = torch.cat(labels, dim=0)
# Compute the AUROC score.
metric = BinaryAUROC()
metric.update(logits, labels)
score = metric.compute().item()
print(f"AUC: {score:.3f}")
def train(model, dataset, device):
dataloader = create_dataloader(dataset, device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
for epoch in range(10):
model.train()
total_loss = 0
########################################################################
# (HIGHLIGHT) Iterate over the dataloader and train the model with all
# mini-batches.
########################################################################
for step, data in enumerate(dataloader):
# Get node pairs with labels for loss calculation.
compacted_seeds = data.compacted_seeds.T
labels = data.labels
# The features of sampled nodes.
x = data.node_features["feat"]
# Forward.
y = model(data.blocks, x)
logits = model.predictor(
y[compacted_seeds[0].long()] * y[compacted_seeds[1].long()]
).squeeze()
# Compute loss.
loss = F.binary_cross_entropy_with_logits(logits, labels.float())
# Backward.
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f"Epoch {epoch:03d} | Loss {total_loss / (step + 1):.3f}")
if __name__ == "__main__":
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Training in {device} mode.")
# Load and preprocess dataset.
print("Loading data...")
dataset = gb.BuiltinDataset("cora").load()
# If a CUDA device is selected, we pin the graph and the features so that
# the GPU can access them.
if device == torch.device("cuda:0"):
dataset.graph.pin_memory_()
dataset.feature.pin_memory_()
in_size = dataset.feature.size("node", None, "feat")[0]
model = GraphSAGE(in_size).to(device)
# Model training.
print("Training...")
train(model, dataset, device)
# Test the model.
print("Testing...")
evaluate(model, dataset, device)
examples/graphbolt/quickstart/node_classification.py
+134
View File
@@ -0,0 +1,134 @@
"""
This example shows how to create a GraphBolt dataloader to sample and train a
node classification model with the Cora dataset.
"""
import dgl.graphbolt as gb
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchmetrics.functional as MF
############################################################################
# (HIGHLIGHT) Create a single process dataloader with dgl graphbolt package.
############################################################################
def create_dataloader(dataset, itemset, device):
# Sample seed nodes from the itemset.
datapipe = gb.ItemSampler(itemset, batch_size=16)
# Copy the mini-batch to the designated device for sampling and training.
datapipe = datapipe.copy_to(device)
# Sample neighbors for the seed nodes.
datapipe = datapipe.sample_neighbor(dataset.graph, fanouts=[4, 2])
# Fetch features for sampled nodes.
datapipe = datapipe.fetch_feature(
dataset.feature, node_feature_keys=["feat"]
)
# Initiate the dataloader for the datapipe.
return gb.DataLoader(datapipe)
class GCN(nn.Module):
def __init__(self, in_size, out_size, hidden_size=16):
super().__init__()
self.layers = nn.ModuleList()
self.layers.append(dglnn.GraphConv(in_size, hidden_size))
self.layers.append(dglnn.GraphConv(hidden_size, out_size))
def forward(self, blocks, x):
hidden_x = x
for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)):
hidden_x = layer(block, hidden_x)
is_last_layer = layer_idx == len(self.layers) - 1
if not is_last_layer:
hidden_x = F.relu(hidden_x)
return hidden_x
@torch.no_grad()
def evaluate(model, dataset, itemset, device):
model.eval()
y = []
y_hats = []
dataloader = create_dataloader(dataset, itemset, device)
for step, data in enumerate(dataloader):
x = data.node_features["feat"]
y.append(data.labels)
y_hats.append(model(data.blocks, x))
return MF.accuracy(
torch.cat(y_hats),
torch.cat(y),
task="multiclass",
num_classes=dataset.tasks[0].metadata["num_classes"],
)
def train(model, dataset, device):
# The first of two tasks in the dataset is node classification.
task = dataset.tasks[0]
dataloader = create_dataloader(dataset, task.train_set, device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
for epoch in range(10):
model.train()
total_loss = 0
########################################################################
# (HIGHLIGHT) Iterate over the dataloader and train the model with all
# mini-batches.
########################################################################
for step, data in enumerate(dataloader):
# The features of sampled nodes.
x = data.node_features["feat"]
# The ground truth labels of the seed nodes.
y = data.labels
# Forward.
y_hat = model(data.blocks, x)
# Compute loss.
loss = F.cross_entropy(y_hat, y)
# Backward.
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
# Evaluate the model.
val_acc = evaluate(model, dataset, task.validation_set, device)
test_acc = evaluate(model, dataset, task.test_set, device)
print(
f"Epoch {epoch:03d} | Loss {total_loss / (step + 1):.3f} | "
f"Val Acc {val_acc.item():.3f} | Test Acc {test_acc.item():.3f}"
)
if __name__ == "__main__":
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Training in {device} mode.")
# Load and preprocess dataset.
print("Loading data...")
dataset = gb.BuiltinDataset("cora").load()
# If a CUDA device is selected, we pin the graph and the features so that
# the GPU can access them.
if device == torch.device("cuda:0"):
dataset.graph.pin_memory_()
dataset.feature.pin_memory_()
in_size = dataset.feature.size("node", None, "feat")[0]
out_size = dataset.tasks[0].metadata["num_classes"]
model = GCN(in_size, out_size).to(device)
# Model training.
print("Training...")
train(model, dataset, device)
examples/graphbolt/rgcn/README.md
+63
View File
@@ -0,0 +1,63 @@
# Node classification on heterogeneous graph with RGCN
This example aims to demonstrate how to run node classification task on heterogeneous graph with **GraphBolt**. Models are not tuned to achieve the best accuracy yet.
## Run on `ogbn-mag` dataset
### 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 | ~5.3GB | 0 | 0GB | ~230s |
| ~1.1GB | ~3GB | 1 | 3.87GB | ~64.6s |
### Accuracies
```
Epoch: 01, Loss: 2.3434, Valid accuracy: 48.23%
Epoch: 02, Loss: 1.5646, Valid accuracy: 48.49%
Epoch: 03, Loss: 1.1633, Valid accuracy: 45.79%
Test accuracy 44.6792
```
## Run on `ogb-lsc-mag240m` dataset
### 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.
> **note:**
`buffer/cache` are highly used during train, it's about 300GB. If more RAM is available, more `buffer/cache` will be consumed as graph size is about 55GB and feature data is about 350GB.
One more thing, first epoch is quite slow as `buffer/cache` is not ready yet. For GPU train, first epoch takes **1030s**.
Even in following epochs, time consumption varies.
| Dataset Size | CPU RAM Usage | Num of GPUs | GPU RAM Usage | Time Per Epoch(Training) |
| ------------ | ------------- | ----------- | ------------- | ------------------------ |
| ~404GB | ~67GB | 0 | 0GB | ~248s |
| ~404GB | ~60GB | 1 | 15GB | ~166s |
### Accuracies
```
Epoch: 01, Loss: 2.1432, Valid accuracy: 50.21%
Epoch: 02, Loss: 1.9267, Valid accuracy: 50.77%
Epoch: 03, Loss: 1.8797, Valid accuracy: 53.38%
```
examples/graphbolt/rgcn/hetero_rgcn.py
+675
View File
@@ -0,0 +1,675 @@
"""
This script is a GraphBolt counterpart of
``/examples/core/rgcn/hetero_rgcn.py``. It demonstrates how to use GraphBolt
to train a R-GCN model for node classification on the Open Graph Benchmark
(OGB) dataset "ogbn-mag" and "ogb-lsc-mag240m". For more details on "ogbn-mag",
please refer to the OGB website: (https://ogb.stanford.edu/docs/linkprop/). For
more details on "ogb-lsc-mag240m", please refer to the OGB website:
(https://ogb.stanford.edu/docs/lsc/mag240m/).
Paper [Modeling Relational Data with Graph Convolutional Networks]
(https://arxiv.org/abs/1703.06103).
This example highlights the user experience of GraphBolt while the model and
training/evaluation procedures are almost identical to the original DGL
implementation. Please refer to original DGL implementation for more details.
This flowchart describes the main functional sequence of the provided example.
main
│
├───> load_dataset
│ │
│ └───> Load dataset
│
├───> rel_graph_embed [HIGHLIGHT]
│ │
│ └───> Generate graph embeddings
│
├───> Instantiate RGCN model
│ │
│ ├───> RelGraphConvLayer (input to hidden)
│ │
│ └───> RelGraphConvLayer (hidden to output)
│
└───> run
│
│
└───> Training loop
│
├───> EntityClassify.forward (RGCN model forward pass)
│
└───> validate and test
│
└───> EntityClassify.evaluate
"""
import argparse
import itertools
import sys
import time
import dgl
import dgl.graphbolt as gb
import dgl.nn as dglnn
import psutil
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import HeteroEmbedding
from ogb.lsc import MAG240MEvaluator
from ogb.nodeproppred import Evaluator
from tqdm import tqdm
def load_dataset(dataset_name):
"""Load the dataset and return the graph, features, train/valid/test sets
and the number of classes.
Here, we use `BuiltInDataset` to load the dataset which returns graph,
features, train/valid/test sets and the number of classes.
"""
dataset = gb.BuiltinDataset(dataset_name).load()
print(f"Loaded dataset: {dataset.tasks[0].metadata['name']}")
graph = dataset.graph
features = dataset.feature
train_set = dataset.tasks[0].train_set
valid_set = dataset.tasks[0].validation_set
test_set = dataset.tasks[0].test_set
num_classes = dataset.tasks[0].metadata["num_classes"]
return (
graph,
features,
train_set,
valid_set,
test_set,
num_classes,
)
def create_dataloader(
name,
graph,
features,
item_set,
device,
batch_size,
fanouts,
shuffle,
num_workers,
):
"""Create a GraphBolt dataloader for training, validation or testing."""
###########################################################################
# Initialize the ItemSampler to sample mini-batches from the dataset.
# `item_set`:
# The set of items to sample from. This is typically the
# training, validation or test set.
# `batch_size`:
# The number of nodes to sample in each mini-batch.
# `shuffle`:
# Whether to shuffle the items in the dataset before sampling.
datapipe = gb.ItemSampler(item_set, batch_size=batch_size, shuffle=shuffle)
# Move the mini-batch to the appropriate device.
# `device`:
# The device to move the mini-batch to.
datapipe = datapipe.copy_to(device)
# Sample neighbors for each seed node in the mini-batch.
# `graph`:
# The graph(FusedCSCSamplingGraph) from which to sample neighbors.
# `fanouts`:
# The number of neighbors to sample for each node in each layer.
datapipe = datapipe.sample_neighbor(
graph,
fanouts=fanouts,
overlap_fetch=args.overlap_graph_fetch,
asynchronous=args.asynchronous,
)
# Fetch the features for each node in the mini-batch.
# `features`:
# The feature store from which to fetch the features.
# `node_feature_keys`:
# The node features to fetch. This is a dictionary where the keys are
# node types and the values are lists of feature names.
node_feature_keys = {"paper": ["feat"]}
if name == "ogb-lsc-mag240m":
node_feature_keys["author"] = ["feat"]
node_feature_keys["institution"] = ["feat"]
datapipe = datapipe.fetch_feature(features, node_feature_keys)
# Create a DataLoader from the datapipe.
# `num_workers`:
# The number of worker processes to use for data loading.
return gb.DataLoader(datapipe, num_workers=num_workers)
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 extract_node_features(name, block, data, node_embed, device):
"""Extract the node features from embedding layer or raw features."""
if name == "ogbn-mag":
input_nodes = {
k: v.to(device) for k, v in block.srcdata[dgl.NID].items()
}
# 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": data.node_features[("paper", "feat")].to(device)}
)
else:
node_features = {
ntype: data.node_features[(ntype, "feat")]
for ntype in block.srctypes
}
# Original feature data are stored in float16 while model weights are
# float32, so we need to convert the features to float32.
node_features = {
k: v.to(device).float() for k, v in node_features.items()
}
return node_features
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 : FusedCSCSamplingGraph
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 = {}
node_type_to_id = graph.node_type_to_id
node_type_offset = graph.node_type_offset
for ntype, ntype_id in node_type_to_id.items():
# Skip the "paper" node type.
if ntype == "paper":
continue
node_num[ntype] = (
node_type_offset[ntype_id + 1] - node_type_offset[ntype_id]
)
print(f"node_num for rel_graph_embed: {node_num}")
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, graph, 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']
etypes = list(graph.edge_type_to_id.keys())
etypes = [gb.etype_str_to_tuple(etype)[1] for etype in etypes]
self.relation_names = etypes
self.relation_names.sort()
self.dropout = 0.5
ntypes = list(graph.node_type_to_id.keys())
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,
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,
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, blocks, h):
for layer, block in zip(self.layers, blocks):
h = layer(block, h)
return h
@torch.no_grad()
def evaluate(
name,
g,
model,
node_embed,
device,
item_set,
features,
num_workers,
):
# Switches the model to evaluation mode.
model.eval()
category = "paper"
# An evaluator for the dataset.
if name == "ogbn-mag":
evaluator = Evaluator(name=name)
else:
evaluator = MAG240MEvaluator()
num_etype = len(g.num_edges)
data_loader = create_dataloader(
name,
g,
features,
item_set,
device,
batch_size=4096,
fanouts=[torch.full((num_etype,), 25), torch.full((num_etype,), 10)],
shuffle=False,
num_workers=num_workers,
)
# To store the predictions.
y_hats = list()
y_true = list()
for data in tqdm(data_loader, desc="Inference"):
# Convert MiniBatch to DGL Blocks and move them to the target device.
blocks = [block.to(device) for block in data.blocks]
node_features = extract_node_features(
name, blocks[0], data, node_embed, device
)
# Generate predictions.
logits = model(blocks, node_features)
logits = logits[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(data.labels[category].long())
y_pred = torch.cat(y_hats, dim=0)
y_true = torch.cat(y_true, dim=0)
y_true = torch.unsqueeze(y_true, 1)
if name == "ogb-lsc-mag240m":
y_pred = y_pred.view(-1)
y_true = y_true.view(-1)
return evaluator.eval({"y_true": y_true, "y_pred": y_pred})["acc"]
def train(
name,
g,
model,
node_embed,
optimizer,
train_set,
valid_set,
device,
features,
num_workers,
num_epochs,
):
print("Start to train...")
category = "paper"
num_etype = len(g.num_edges)
data_loader = create_dataloader(
name,
g,
features,
train_set,
device,
batch_size=1024,
fanouts=[torch.full((num_etype,), 25), torch.full((num_etype,), 10)],
shuffle=True,
num_workers=num_workers,
)
# 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(num_epochs):
num_train = len(train_set)
t0 = time.time()
model.train()
total_loss = 0
for data in tqdm(data_loader, desc=f"Training~Epoch {epoch + 1:02d}"):
# Convert MiniBatch to DGL Blocks and move them to the target
# device.
blocks = [block.to(device) for block in data.blocks]
# Fetch the number of seed nodes in the batch.
num_seeds = blocks[-1].num_dst_nodes(category)
# Extract the node features from embedding layer or raw features.
node_features = extract_node_features(
name, blocks[0], data, node_embed, device
)
# Reset gradients.
optimizer.zero_grad()
# Generate predictions.
logits = model(blocks, node_features)[category]
y_hat = logits.log_softmax(dim=-1)
loss = F.nll_loss(y_hat, data.labels[category].long())
loss.backward()
optimizer.step()
total_loss += loss.item() * num_seeds
t1 = time.time()
loss = total_loss / num_train
# Evaluate the model on the val/test set.
print("Evaluating the model on the validation set.")
valid_acc = evaluate(
name, g, model, node_embed, device, valid_set, features, num_workers
)
print("Finish evaluating on validation set.")
print(
f"Epoch: {epoch + 1:02d}, "
f"Loss: {loss:.4f}, "
f"Valid accuracy: {100 * valid_acc:.2f}%, "
f"Time {t1 - t0:.4f}"
)
def main(args):
device = torch.device(
"cuda" if args.num_gpus > 0 and torch.cuda.is_available() else "cpu"
)
# Load dataset.
(
g,
features,
train_set,
valid_set,
test_set,
num_classes,
) = load_dataset(args.dataset)
# Move the dataset to the pinned memory to enable GPU access.
args.overlap_graph_fetch = False
args.asynchronous = False
if device == torch.device("cuda"):
g = g.pin_memory_()
features = features.pin_memory_()
# Enable optimizations for sampling on the GPU.
args.overlap_graph_fetch = True
args.asynchronous = True
feat_size = features.size("node", "paper", "feat")[0]
# As `ogb-lsc-mag240m` is a large dataset, features of `author` and
# `institution` are generated in advance and stored in the feature store.
# For `ogbn-mag`, we generate the features on the fly.
embed_layer = None
if args.dataset == "ogbn-mag":
# Create the embedding layer and move it to the appropriate device.
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())}"
)
if embed_layer is not None:
embed_layer.reset_parameters()
model.reset_parameters()
# `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 = 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,
model,
embed_layer,
optimizer,
train_set,
valid_set,
device,
features,
args.num_workers,
args.num_epochs,
)
print("Testing...")
test_acc = evaluate(
args.dataset,
g,
model,
embed_layer,
device,
test_set,
features,
args.num_workers,
)
print(f"Test accuracy {test_acc*100:.4f}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="GraphBolt RGCN")
parser.add_argument(
"--dataset",
type=str,
default="ogbn-mag",
choices=["ogbn-mag", "ogb-lsc-mag240m"],
help="Dataset name. Possible values: ogbn-mag, ogb-lsc-mag240m",
)
parser.add_argument("--num_epochs", type=int, default=3)
parser.add_argument("--num_workers", type=int, default=0)
parser.add_argument("--num_gpus", type=int, default=1)
args = parser.parse_args()
main(args)
examples/graphbolt/sparse/graphsage.py
+266
View File
@@ -0,0 +1,266 @@
"""
This script demonstrate how to use dgl sparse library to sample on graph and
train model. It trains and tests a GraphSAGE model using the sparse sample and
compact operators to sample submatrix from the whole matrix.
This flowchart describes the main functional sequence of the provided example.
main
│
├───> Load and preprocess full dataset
│
├───> Instantiate SAGE model
│
├───> train
│ │
│ └───> Training loop
│ │
│ ├───> Sample submatrix
│ │
│ └───> SAGE.forward
└───> test
│
├───> Sample submatrix
│
└───> Evaluate the model
"""
import argparse
from functools import partial
import dgl.graphbolt as gb
import dgl.sparse as dglsp
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchmetrics.functional as MF
from dgl.graphbolt.subgraph_sampler import SubgraphSampler
from torch.utils.data import functional_datapipe
from tqdm import tqdm
class SAGEConv(nn.Module):
r"""GraphSAGE layer from `Inductive Representation Learning on
Large Graphs <https://arxiv.org/pdf/1706.02216.pdf>`__
"""
def __init__(
self,
in_feats,
out_feats,
):
super(SAGEConv, self).__init__()
self._in_src_feats, self._in_dst_feats = in_feats, in_feats
self._out_feats = out_feats
self.fc_neigh = nn.Linear(self._in_src_feats, out_feats, bias=False)
self.fc_self = nn.Linear(self._in_dst_feats, out_feats, bias=True)
self.reset_parameters()
def reset_parameters(self):
gain = nn.init.calculate_gain("relu")
nn.init.xavier_uniform_(self.fc_self.weight, gain=gain)
nn.init.xavier_uniform_(self.fc_neigh.weight, gain=gain)
def forward(self, A, feat):
feat_src = feat
feat_dst = feat[: A.shape[1]]
# Aggregator type: mean.
srcdata = self.fc_neigh(feat_src)
# Divided by degree.
D_hat = dglsp.diag(A.sum(0)) ** -1
A_div = A @ D_hat
# Conv neighbors.
dstdata = A_div.T @ srcdata
rst = self.fc_self(feat_dst) + dstdata
return rst
class SAGE(nn.Module):
def __init__(self, in_size, hid_size, out_size):
super().__init__()
self.layers = nn.ModuleList()
# Three-layer GraphSAGE-gcn.
self.layers.append(SAGEConv(in_size, hid_size))
self.layers.append(SAGEConv(hid_size, hid_size))
self.layers.append(SAGEConv(hid_size, out_size))
self.dropout = nn.Dropout(0.5)
self.hid_size = hid_size
self.out_size = out_size
def forward(self, sampled_matrices, x):
hidden_x = x
for layer_idx, (layer, sampled_matrix) in enumerate(
zip(self.layers, sampled_matrices)
):
hidden_x = layer(sampled_matrix, hidden_x)
if layer_idx != len(self.layers) - 1:
hidden_x = F.relu(hidden_x)
hidden_x = self.dropout(hidden_x)
return hidden_x
@functional_datapipe("sample_sparse_neighbor")
class SparseNeighborSampler(SubgraphSampler):
def __init__(self, datapipe, matrix, fanouts):
super().__init__(datapipe)
self.matrix = matrix
# Convert fanouts to a list of tensors.
self.fanouts = []
for fanout in fanouts:
if not isinstance(fanout, torch.Tensor):
fanout = torch.LongTensor([int(fanout)])
self.fanouts.insert(0, fanout)
def sample_subgraphs(self, seeds, seeds_timestamp=None):
sampled_matrices = []
src = seeds.long()
#####################################################################
# (HIGHLIGHT) Using the sparse sample operator to preform random
# sampling on the neighboring nodes of the seeds nodes. The sparse
# compact operator is then employed to compact and relabel the sampled
# matrix, resulting in the sampled matrix and the relabel index.
#####################################################################
for fanout in self.fanouts:
# Sample neighbors.
sampled_matrix = self.matrix.sample(1, fanout, ids=src).coalesce()
# Compact the sampled matrix.
compacted_mat, row_ids = sampled_matrix.compact(0)
sampled_matrices.insert(0, compacted_mat)
src = row_ids
return src, sampled_matrices
############################################################################
# (HIGHLIGHT) Create a multi-process dataloader with dgl graphbolt package.
############################################################################
def create_dataloader(A, fanouts, ids, features, device):
datapipe = gb.ItemSampler(ids, batch_size=1024)
# Customize graphbolt sampler by sparse.
datapipe = datapipe.sample_sparse_neighbor(A, fanouts)
# Use grapbolt to fetch features.
datapipe = datapipe.fetch_feature(features, node_feature_keys=["feat"])
datapipe = datapipe.copy_to(device)
dataloader = gb.DataLoader(datapipe)
return dataloader
def evaluate(model, dataloader, num_classes):
model.eval()
ys = []
y_hats = []
for it, data in tqdm(enumerate(dataloader), "Evaluating"):
with torch.no_grad():
node_feature = data.node_features["feat"].float()
blocks = data.sampled_subgraphs
y = data.labels
ys.append(y)
y_hats.append(model(blocks, node_feature))
return MF.accuracy(
torch.cat(y_hats),
torch.cat(ys),
task="multiclass",
num_classes=num_classes,
)
def validate(device, dataset, model, num_classes):
test_set = dataset.tasks[0].test_set
test_dataloader = create_dataloader(
A, [10, 10, 10], test_set, features, device
)
acc = evaluate(model, test_dataloader, num_classes)
return acc
def train(device, A, features, dataset, num_classes, model):
# Create sampler & dataloader.
train_set = dataset.tasks[0].train_set
train_dataloader = create_dataloader(
A, [10, 10, 10], train_set, features, device
)
valid_set = dataset.tasks[0].validation_set
val_dataloader = create_dataloader(
A, [10, 10, 10], valid_set, features, device
)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4)
for epoch in range(10):
model.train()
total_loss = 0
for it, data in tqdm(enumerate(train_dataloader), "Training"):
node_feature = data.node_features["feat"].float()
blocks = data.sampled_subgraphs
y = data.labels
y_hat = model(blocks, node_feature)
loss = F.cross_entropy(y_hat, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
acc = evaluate(model, val_dataloader, num_classes)
print(
"Epoch {:05d} | Loss {:.4f} | Accuracy {:.4f} ".format(
epoch, total_loss / (it + 1), acc.item()
)
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="GraphSAGE")
parser.add_argument(
"--mode",
default="gpu",
choices=["cpu", "gpu"],
help="Training mode. 'cpu' for CPU training, 'gpu' for GPU training.",
)
args = parser.parse_args()
if not torch.cuda.is_available():
args.mode = "cpu"
print(f"Training in {args.mode} mode.")
#####################################################################
# (HIGHLIGHT) This example implements a graphSAGE algorithm by sparse
# operators, which involves sampling a subgraph from a full graph and
# conducting training.
#
# First, the whole graph is loaded onto the CPU or GPU and transformed
# to sparse matrix. To obtain the training subgraph, it samples three
# submatrices by seed nodes, which contains their randomly sampled
# 1-hop, 2-hop, and 3-hop neighbors. Then, the features of the
# subgraph are input to the network for training.
#####################################################################
# Load and preprocess dataset.
print("Loading data")
device = torch.device("cpu" if args.mode == "cpu" else "cuda")
dataset = gb.BuiltinDataset("ogbn-products").load()
g = dataset.graph
features = dataset.feature
# Create GraphSAGE model.
in_size = features.size("node", None, "feat")[0]
num_classes = dataset.tasks[0].metadata["num_classes"]
out_size = num_classes
model = SAGE(in_size, 256, out_size).to(device)
# Create sparse.
N = g.num_nodes
A = dglsp.from_csc(g.csc_indptr.long(), g.indices.long(), shape=(N, N))
# Model training.
print("Training...")
train(device, A, features, dataset, num_classes, model)
# Test the model.
print("Testing...")
acc = validate(device, dataset, model, num_classes)
print(f"Test accuracy {acc:.4f}")
examples/graphbolt/temporal_link_prediction.py
+333
View File
@@ -0,0 +1,333 @@
"""
This script trains and tests a Heterogeneous GraphSAGE model for link
prediction with temporal information using graphbolt dataloader.
While node classification predicts labels for nodes based on their
local neighborhoods, link prediction assesses the likelihood of an edge
existing between two nodes, necessitating different sampling strategies
that account for pairs of nodes and their joint neighborhoods.
An additional temporal attribute is provided in both graph and TVT sets,
ensuring that during sampling, only neighbors whose timestamps are earlier
than the seed timestamp will be sampled.
This flowchart describes the main functional sequence of the provided example.
main
│
├───> OnDiskDataset pre-processing
│
├───> Instantiate HeteroSAGE model
│
├───> train
│ │
│ ├───> Get graphbolt dataloader (HIGHLIGHT)
│ │
│ └───> Training loop
│ │
│ ├───> HeteroSAGE.forward
│ │
│ └───> Validation set evaluation
│
└───> Test set evaluation
"""
import argparse
import os
import time
import dgl.graphbolt as gb
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
import tqdm
from dgl.data.utils import download, extract_archive
TIMESTAMP_FEATURE_NAME = "__timestamp__"
NODE_FEATURE_KEYS = {
"Product": ["categoryId"],
"Query": ["categoryId"],
}
TARGET_TYPE = ("Query", "Click", "Product")
ALL_TYPES = [
TARGET_TYPE,
("Product", "reverse_Click", "Query"),
("Product", "reverse_QueryResult", "Query"),
("Query", "QueryResult", "Product"),
]
class CategoricalEncoder(nn.Module):
def __init__(
self,
num_categories,
out_size,
):
super().__init__()
self.embed = nn.Embedding(num_categories, out_size)
self.reset_parameters()
def reset_parameters(self):
nn.init.xavier_uniform_(self.embed.weight)
def forward(self, input_feat: torch.Tensor):
return self.embed(input_feat.view(-1))
class HeteroSAGE(nn.Module):
def __init__(self, in_size, hidden_size):
super().__init__()
self.layers = nn.ModuleList()
sizes = [in_size, hidden_size]
for size in sizes:
self.layers.append(
dglnn.HeteroGraphConv(
{
etype: dglnn.SAGEConv(
size,
hidden_size,
"mean",
)
for etype in ALL_TYPES
},
aggregate="sum",
)
)
self.predictor = nn.Sequential(
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, 1),
)
def forward(self, blocks, X_node_dict):
H_node_dict = X_node_dict
for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)):
H_node_dict = layer(block, H_node_dict)
is_last_layer = layer_idx == len(self.layers) - 1
if not is_last_layer:
H_node_dict = {
ntype: F.relu(H) for ntype, H in H_node_dict.items()
}
return H_node_dict
def create_dataloader(args, graph, features, itemset, is_train=True):
datapipe = gb.ItemSampler(
itemset,
batch_size=args.train_batch_size if is_train else args.eval_batch_size,
shuffle=is_train,
)
if args.storage_device != "cpu":
datapipe = datapipe.copy_to(device=args.device)
############################################################################
# [Input]:
# 'datapipe' is either 'ItemSampler' or 'UniformNegativeSampler' depending
# on whether training is needed ('is_train'),
# 'graph': The network topology for sampling.
# 'args.fanout': Number of neighbors to sample per node.
# [Output]:
# A NeighborSampler object to sample neighbors.
# [Role]:
# Initialize a neighbor sampler for sampling the neighborhoods of nodes with
# considering of temporal information. Only neighbors that is earlier than
# the seed will be sampled.
############################################################################
datapipe = getattr(datapipe, args.sample_mode)(
graph,
args.fanout if is_train else [-1],
node_timestamp_attr_name=TIMESTAMP_FEATURE_NAME,
edge_timestamp_attr_name=TIMESTAMP_FEATURE_NAME,
)
datapipe = datapipe.fetch_feature(
features, node_feature_keys=NODE_FEATURE_KEYS
)
if args.storage_device == "cpu":
datapipe = datapipe.copy_to(device=args.device)
dataloader = gb.DataLoader(
datapipe,
num_workers=args.num_workers,
)
# Return the fully-initialized DataLoader object.
return dataloader
def train(args, model, graph, features, train_set, encoders):
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
dataloader = create_dataloader(args, graph, features, train_set)
for epoch in range(args.epochs):
model.train()
total_loss = 0
start_epoch_time = time.time()
for step, data in tqdm.tqdm(enumerate(dataloader)):
# Get node pairs with labels for loss calculation.
compacted_seeds = data.compacted_seeds[
gb.etype_tuple_to_str(TARGET_TYPE)
].T
labels = data.labels
node_feature = {}
for ntype, keys in NODE_FEATURE_KEYS.items():
ntype, feat = ntype, keys[0]
node_feature[ntype] = data.node_features[
(ntype, feat)
].squeeze()
blocks = data.blocks
# Get the embeddings of the input nodes.
X_node_dict = {
ntype: encoders[ntype](feat)
for ntype, feat in node_feature.items()
}
X_node_dict = model(blocks, X_node_dict)
src_type, _, dst_type = TARGET_TYPE
logits = model.predictor(
X_node_dict[src_type][compacted_seeds[0]]
* X_node_dict[dst_type][compacted_seeds[1]]
).squeeze()
# Compute loss.
loss = F.binary_cross_entropy_with_logits(
logits, labels[gb.etype_tuple_to_str(TARGET_TYPE)].float()
)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
if step + 1 == args.early_stop:
# Early stopping requires a new dataloader to reset its state.
dataloader = create_dataloader(args, graph, features, train_set)
break
end_epoch_time = time.time()
print(
f"Epoch {epoch:05d} | "
f"Loss {(total_loss) / (step + 1):.4f} | "
f"Time {(end_epoch_time - start_epoch_time):.4f} s"
)
def parse_args():
parser = argparse.ArgumentParser(description="diginetica-r2ne (GraphBolt)")
parser.add_argument("--epochs", type=int, default=10)
parser.add_argument("--lr", type=float, default=0.0005)
parser.add_argument("--neg-ratio", type=int, default=1)
parser.add_argument("--train-batch-size", type=int, default=1024)
parser.add_argument("--eval-batch-size", type=int, default=1024)
parser.add_argument("--num-workers", type=int, default=0)
parser.add_argument(
"--dataset",
default="diginetica-r2ne",
choices=["diginetica-r2ne"],
help="Dataset.",
)
parser.add_argument(
"--early-stop",
type=int,
default=0,
help="0 means no early stop, otherwise stop at the input-th step",
)
parser.add_argument(
"--fanout",
type=str,
default="20,20",
help="Fan-out of neighbor sampling. Default: 20, 20",
)
parser.add_argument(
"--exclude-edges",
type=int,
default=1,
help="Whether to exclude reverse edges during sampling. Default: 1",
)
parser.add_argument(
"--mode",
default="cpu-cuda",
choices=["cpu-cpu", "cpu-cuda", "cuda-cuda"],
help="Dataset storage placement and Train device: 'cpu' for CPU and RAM,"
" 'pinned' for pinned memory in RAM, 'cuda' for GPU and GPU memory.",
)
parser.add_argument(
"--sample-mode",
default="temporal_sample_neighbor",
choices=["temporal_sample_neighbor", "temporal_sample_layer_neighbor"],
help="The sampling function when doing layerwise sampling.",
)
return parser.parse_args()
def download_datasets(name, root="datasets"):
url = "https://dgl-data.s3-accelerate.amazonaws.com/dataset/"
dataset_dir = os.path.join(root, name)
if not os.path.exists(dataset_dir):
url += name + ".zip"
os.makedirs(root, exist_ok=True)
zip_file_path = os.path.join(root, name + ".zip")
download(url, path=zip_file_path)
extract_archive(zip_file_path, root, overwrite=True)
os.remove(zip_file_path)
return dataset_dir
def main(args):
if not torch.cuda.is_available():
args.mode = "cpu-cpu"
print(f"Training in {args.mode} mode.")
args.storage_device, args.device = args.mode.split("-")
args.device = torch.device(args.device)
# Load and preprocess dataset.
print("Loading data")
# TODO: Add the datasets to built-in.
dataset_path = download_datasets(args.dataset)
dataset = gb.OnDiskDataset(dataset_path).load()
# Move the dataset to the selected storage.
graph = dataset.graph.to(args.storage_device)
features = dataset.feature.to(args.storage_device)
train_set = dataset.tasks[0].train_set
# TODO: Fix the dataset so that this modification is not needed. node_pairs
# needs to be cast into graph.indices.dtype, which is int32.
train_set._itemsets["Query:Click:Product"]._items = tuple(
item.to(graph.indices.dtype if i == 0 else None)
for i, item in enumerate(
train_set._itemsets["Query:Click:Product"]._items
)
)
args.fanout = list(map(int, args.fanout.split(",")))
in_size = 128
hidden_channels = 256
query_size = features.metadata("node", "Query", "categoryId")[
"num_categories"
]
product_size = features.metadata("node", "Product", "categoryId")[
"num_categories"
]
args.device = torch.device(args.device)
model = HeteroSAGE(in_size, hidden_channels).to(args.device)
encoders = {
"Query": CategoricalEncoder(query_size, in_size).to(args.device),
"Product": CategoricalEncoder(product_size, in_size).to(args.device),
}
# Model training.
print("Training...")
train(args, model, graph, features, train_set, encoders)
if __name__ == "__main__":
args = parse_args()
main(args)