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/distributed/graphsage/README.md
+221
View File
@@ -0,0 +1,221 @@
## Distributed training
This is an example of training GraphSage in a distributed fashion. Before training, please install some python libs by pip:
```
pip3 install ogb
```
**Requires PyTorch 1.12.0+ to work.**
To train GraphSage, it has five steps:
### Step 0: Setup a Distributed File System
* You may skip this step if your cluster already has folder(s) synchronized across machines.
To perform distributed training, files and codes need to be accessed across multiple machines. A distributed file system would perfectly handle the job (i.e., NFS, Ceph).
#### Server side setup
Here is an example of how to setup NFS. First, install essential libs on the storage server
```
sudo apt-get install nfs-kernel-server
```
Below we assume the user account is `ubuntu` and we create a directory of `workspace` in the home directory.
```
mkdir -p /home/ubuntu/workspace
```
We assume that the all servers are under a subnet with ip range `192.168.0.0` to `192.168.255.255`. The exports configuration needs to be modifed to
```
sudo vim /etc/exports
# add the following line
/home/ubuntu/workspace 192.168.0.0/16(rw,sync,no_subtree_check)
```
The server's internal ip can be checked via `ifconfig` or `ip`. If the ip does not begin with `192.168`, then you may use
```
/home/ubuntu/workspace 10.0.0.0/8(rw,sync,no_subtree_check)
/home/ubuntu/workspace 172.16.0.0/12(rw,sync,no_subtree_check)
```
Then restart NFS, the setup on server side is finished.
```
sudo systemctl restart nfs-kernel-server
```
For configraution details, please refer to [NFS ArchWiki](https://wiki.archlinux.org/index.php/NFS).
#### Client side setup
To use NFS, clients also require to install essential packages
```
sudo apt-get install nfs-common
```
You can either mount the NFS manually
```
mkdir -p /home/ubuntu/workspace
sudo mount -t nfs <nfs-server-ip>:/home/ubuntu/workspace /home/ubuntu/workspace
```
or edit the fstab so the folder will be mounted automatically
```
# vim /etc/fstab
## append the following line to the file
<nfs-server-ip>:/home/ubuntu/workspace /home/ubuntu/workspace nfs defaults 0 0
```
Then run `mount -a`.
Now go to `/home/ubuntu/workspace` and clone the DGL Github repository.
### Step 1: set IP configuration file.
User need to set their own IP configuration file `ip_config.txt` before training. For example, if we have four machines in current cluster, the IP configuration
could like this:
```
172.31.19.1
172.31.23.205
172.31.29.175
172.31.16.98
```
Users need to make sure that the master node (node-0) has right permission to ssh to all the other nodes without password authentication.
[This link](https://linuxize.com/post/how-to-setup-passwordless-ssh-login/) provides instructions of setting passwordless SSH login.
### Step 2: partition the graph.
The example provides a script to partition some builtin graphs such as Reddit and OGB product graph.
If we want to train GraphSage on 4 machines, we need to partition the graph into 4 parts.
In this example, we partition the ogbn-products graph into 4 parts with Metis on node-0. The partitions are balanced with respect to
the number of nodes, the number of edges and the number of labelled nodes.
```
python3 partition_graph.py --dataset ogbn-products --num_parts 4 --balance_train --balance_edges
```
This script generates partitioned graphs and store them in the directory called `data`.
### Step 3: Launch distributed jobs
DGL provides a script to launch the training job in the cluster. `part_config` and `ip_config`
specify relative paths to the path of the workspace.
The command below launches one process per machine for both sampling and training.
```
python3 ~/workspace/dgl/tools/launch.py \
--workspace ~/workspace/dgl/examples/distributed/graphsage/ \
--num_trainers 1 \
--num_samplers 0 \
--num_servers 1 \
--part_config data/ogbn-products.json \
--ip_config ip_config.txt \
"python3 node_classification.py --graph_name ogbn-products --ip_config ip_config.txt --num_epochs 30 --batch_size 1000"
```
By default, this code will run on CPU. If you have GPU support, you can just add a `--num_gpus` argument in user command:
```
python3 ~/workspace/dgl/tools/launch.py \
--workspace ~/workspace/dgl/examples/distributed/graphsage/ \
--num_trainers 4 \
--num_samplers 0 \
--num_servers 1 \
--part_config data/ogbn-products.json \
--ip_config ip_config.txt \
"python3 node_classification.py --graph_name ogbn-products --ip_config ip_config.txt --num_epochs 30 --batch_size 1000 --num_gpus 4"
```
Unsupervised training(train with link prediction dataloader).
```
python3 ~/workspace/dgl/tools/launch.py \
--workspace ~/workspace/dgl/examples/distributed/graphsage/ \
--num_trainers 1 \
--num_samplers 0 \
--num_servers 1 \
--part_config data/ogbn-products.json \
--ip_config ip_config.txt \
"python3 node_classification_unsupervised.py --graph_name ogbn-products --ip_config ip_config.txt --num_epochs 30 --batch_size 1000 --remove_edge"
```
### Running with GraphBolt
In order to run with `GraphBolt`, we need to partition graph into `GraphBolt` data formats.Please note that both `DGL` and `GraphBolt` partitions are saved together.
If we have already partitioned into `DGL` format, just convert them directly like below:
```
python3 -c "import dgl; dgl.distributed.dgl_partition_to_graphbolt('ogbn-products.json')"
```
Or partition from scratch like this:
```
python3 partition_graph.py --dataset ogbn-products --num_parts 2 --balance_train --balance_edges --use_graphbolt
```
#### Partition sizes compared to DGL
Compared to `DGL`, `GraphBolt` partitions are much smaller(reduced to **16%** and **19%** for `ogbn-products` and `ogbn-papers100M` respectively).
`ogbn-products`
| Data Formats | File Name | Part 0 | Part 1 |
| ------------ | ---------------------------- | ------ | ------ |
| DGL | graph.dgl | 1.5GB | 1.6GB |
| GraphBolt | fused_csc_sampling_graph.pt | 255MB | 265MB |
`ogbn-papers100M`
| Data Formats | File Name | Part 0 | Part 1 |
| ------------ | ---------------------------- | ------ | ------ |
| DGL | graph.dgl | 23GB | 22GB |
| GraphBolt | fused_csc_sampling_graph.pt | 4.4GB | 4.1GB |
Then run example with `--use_graphbolt`.
```
python3 ~/workspace/dgl/tools/launch.py \
--workspace ~/workspace/dgl/examples/distributed/graphsage/ \
--num_trainers 4 \
--num_samplers 0 \
--num_servers 2 \
--part_config data/ogbn-products.json \
--ip_config ip_config.txt \
"python3 node_classification.py --graph_name ogbn-products --ip_config ip_config.txt --num_epochs 10 --use_graphbolt"
```
#### Performance compared to `DGL`
Compared to `DGL`, `GraphBolt`'s sampler works faster(reduced to **80%** and **77%** for `ogbn-products` and `ogbn-papers100M` respectively). `Min` and `Max` are statistics of all trainers on all nodes(machines).
As for RAM usage, the shared memory(measured by **shared** field of `free` command) usage is decreased due to smaller graph partitions in `GraphBolt` though the peak memory used by processes(measured by **used** field of `free` command) does not decrease.
`ogbn-products`
| Data Formats | Sample Time Per Epoch (CPU) | Test Accuracy (10 epochs) | shared | used (peak) |
| ------------ | --------------------------- | -------------------------------- | ----- | ---- |
| DGL | Min: 1.2884s, Max: 1.4159s | Min: 64.38%, Max: 70.42% | 2.4GB | 7.8GB|
| GraphBolt | Min: 1.0589s, Max: 1.1400s | Min: 61.68%, Max: 71.23% | 1.1GB | 7.8GB|
`ogbn-papers100M`
| Data Formats | Sample Time Per Epoch (CPU) | Test Accuracy (10 epochs) | shared | used (peak) |
| ------------ | --------------------------- | -------------------------------- | ----- | ---- |
| DGL | Min: 5.5570s, Max: 6.1900s | Min: 29.12%, Max: 34.33% | 84GB | 43GB |
| GraphBolt | Min: 4.5046s, Max: 4.7718s | Min: 29.11%, Max: 33.49% | 67GB | 43GB |
examples/distributed/graphsage/node_classification.py
+468
View File
@@ -0,0 +1,468 @@
import argparse
import socket
import time
import dgl
import dgl.distributed
import dgl.nn.pytorch as dglnn
import numpy as np
import torch as th
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import tqdm
class DistSAGE(nn.Module):
"""
SAGE model for distributed train and evaluation.
Parameters
----------
in_feats : int
Feature dimension.
n_hidden : int
Hidden layer dimension.
n_classes : int
Number of classes.
n_layers : int
Number of layers.
activation : callable
Activation function.
dropout : float
Dropout value.
"""
def __init__(
self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
):
super().__init__()
self.n_layers = n_layers
self.n_hidden = n_hidden
self.n_classes = n_classes
self.layers = nn.ModuleList()
self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, "mean"))
for _ in range(1, n_layers - 1):
self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, "mean"))
self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, "mean"))
self.dropout = nn.Dropout(dropout)
self.activation = activation
def forward(self, blocks, x):
"""
Forward function.
Parameters
----------
blocks : List[DGLBlock]
Sampled blocks.
x : DistTensor
Feature data.
"""
h = x
for i, (layer, block) in enumerate(zip(self.layers, blocks)):
h = layer(block, h)
if i != len(self.layers) - 1:
h = self.activation(h)
h = self.dropout(h)
return h
def inference(self, g, x, batch_size, device):
"""
Distributed layer-wise inference with the GraphSAGE model on full
neighbors.
Parameters
----------
g : DistGraph
Input Graph for inference.
x : DistTensor
Node feature data of input graph.
Returns
-------
DistTensor
Inference results.
"""
# Split nodes to each trainer.
nodes = dgl.distributed.node_split(
np.arange(g.num_nodes()),
g.get_partition_book(),
force_even=True,
)
for i, layer in enumerate(self.layers):
# Create DistTensor to save forward results.
if i == len(self.layers) - 1:
out_dim = self.n_classes
name = "h_last"
else:
out_dim = self.n_hidden
name = "h"
y = dgl.distributed.DistTensor(
(g.num_nodes(), out_dim),
th.float32,
name,
persistent=True,
)
print(f"|V|={g.num_nodes()}, inference batch size: {batch_size}")
# `-1` indicates all inbound edges will be inlcuded, namely, full
# neighbor sampling.
sampler = dgl.dataloading.NeighborSampler([-1])
dataloader = dgl.distributed.DistNodeDataLoader(
g,
nodes,
sampler,
batch_size=batch_size,
shuffle=False,
drop_last=False,
)
for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
block = blocks[0].to(device)
h = x[input_nodes].to(device)
h_dst = h[: block.number_of_dst_nodes()]
h = layer(block, (h, h_dst))
if i != len(self.layers) - 1:
h = self.activation(h)
h = self.dropout(h)
# Copy back to CPU as DistTensor requires data reside on CPU.
y[output_nodes] = h.cpu()
x = y
# Synchronize trainers.
g.barrier()
return x
def compute_acc(pred, labels):
"""
Compute the accuracy of prediction given the labels.
Parameters
----------
pred : torch.Tensor
Predicted labels.
labels : torch.Tensor
Ground-truth labels.
Returns
-------
float
Accuracy.
"""
labels = labels.long()
return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred)
def evaluate(model, g, inputs, labels, val_nid, test_nid, batch_size, device):
"""
Evaluate the model on the validation and test set.
Parameters
----------
model : DistSAGE
The model to be evaluated.
g : DistGraph
The entire graph.
inputs : DistTensor
The feature data of all the nodes.
labels : DistTensor
The labels of all the nodes.
val_nid : torch.Tensor
The node IDs for validation.
test_nid : torch.Tensor
The node IDs for test.
batch_size : int
Batch size for evaluation.
device : torch.Device
The target device to evaluate on.
Returns
-------
float
Validation accuracy.
float
Test accuracy.
"""
model.eval()
with th.no_grad():
pred = model.inference(g, inputs, batch_size, device)
model.train()
return compute_acc(pred[val_nid], labels[val_nid]), compute_acc(
pred[test_nid], labels[test_nid]
)
def run(args, device, data):
"""
Train and evaluate DistSAGE.
Parameters
----------
args : argparse.Args
Arguments for train and evaluate.
device : torch.Device
Target device for train and evaluate.
data : Packed Data
Packed data includes train/val/test IDs, feature dimension,
number of classes, graph.
"""
train_nid, val_nid, test_nid, in_feats, n_classes, g = data
sampler = dgl.dataloading.NeighborSampler(
[int(fanout) for fanout in args.fan_out.split(",")]
)
dataloader = dgl.distributed.DistNodeDataLoader(
g,
train_nid,
sampler,
batch_size=args.batch_size,
shuffle=True,
drop_last=False,
)
model = DistSAGE(
in_feats,
args.num_hidden,
n_classes,
args.num_layers,
F.relu,
args.dropout,
)
model = model.to(device)
if args.num_gpus == 0:
model = th.nn.parallel.DistributedDataParallel(model)
else:
model = th.nn.parallel.DistributedDataParallel(
model, device_ids=[device], output_device=device
)
loss_fcn = nn.CrossEntropyLoss()
loss_fcn = loss_fcn.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
# Training loop.
iter_tput = []
epoch = 0
epoch_time = []
test_acc = 0.0
for _ in range(args.num_epochs):
epoch += 1
tic = time.time()
# Various time statistics.
sample_time = 0
forward_time = 0
backward_time = 0
update_time = 0
num_seeds = 0
num_inputs = 0
start = time.time()
step_time = []
with model.join():
for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
tic_step = time.time()
sample_time += tic_step - start
# Slice feature and label.
batch_inputs = g.ndata["features"][input_nodes]
batch_labels = g.ndata["labels"][seeds].long()
num_seeds += len(blocks[-1].dstdata[dgl.NID])
num_inputs += len(blocks[0].srcdata[dgl.NID])
# Move to target device.
blocks = [block.to(device) for block in blocks]
batch_inputs = batch_inputs.to(device)
batch_labels = batch_labels.to(device)
# Compute loss and prediction.
start = time.time()
batch_pred = model(blocks, batch_inputs)
loss = loss_fcn(batch_pred, batch_labels)
forward_end = time.time()
optimizer.zero_grad()
loss.backward()
compute_end = time.time()
forward_time += forward_end - start
backward_time += compute_end - forward_end
optimizer.step()
update_time += time.time() - compute_end
step_t = time.time() - tic_step
step_time.append(step_t)
iter_tput.append(len(blocks[-1].dstdata[dgl.NID]) / step_t)
if (step + 1) % args.log_every == 0:
acc = compute_acc(batch_pred, batch_labels)
gpu_mem_alloc = (
th.cuda.max_memory_allocated() / 1000000
if th.cuda.is_available()
else 0
)
sample_speed = np.mean(iter_tput[-args.log_every :])
mean_step_time = np.mean(step_time[-args.log_every :])
print(
f"Part {g.rank()} | Epoch {epoch:05d} | Step {step:05d}"
f" | Loss {loss.item():.4f} | Train Acc {acc.item():.4f}"
f" | Speed (samples/sec) {sample_speed:.4f}"
f" | GPU {gpu_mem_alloc:.1f} MB | "
f"Mean step time {mean_step_time:.3f} s"
)
start = time.time()
toc = time.time()
print(
f"Part {g.rank()}, Epoch Time(s): {toc - tic:.4f}, "
f"sample+data_copy: {sample_time:.4f}, forward: {forward_time:.4f},"
f" backward: {backward_time:.4f}, update: {update_time:.4f}, "
f"#seeds: {num_seeds}, #inputs: {num_inputs}"
)
epoch_time.append(toc - tic)
if epoch % args.eval_every == 0 or epoch == args.num_epochs:
start = time.time()
val_acc, test_acc = evaluate(
model.module,
g,
g.ndata["features"],
g.ndata["labels"],
val_nid,
test_nid,
args.batch_size_eval,
device,
)
print(
f"Part {g.rank()}, Val Acc {val_acc:.4f}, "
f"Test Acc {test_acc:.4f}, time: {time.time() - start:.4f}"
)
return np.mean(epoch_time[-int(args.num_epochs * 0.8) :]), test_acc
def main(args):
"""
Main function.
"""
host_name = socket.gethostname()
print(f"{host_name}: Initializing DistDGL.")
dgl.distributed.initialize(args.ip_config, use_graphbolt=args.use_graphbolt)
print(f"{host_name}: Initializing PyTorch process group.")
th.distributed.init_process_group(backend=args.backend)
print(f"{host_name}: Initializing DistGraph.")
g = dgl.distributed.DistGraph(args.graph_name, part_config=args.part_config)
print(f"Rank of {host_name}: {g.rank()}")
# Split train/val/test IDs for each trainer.
pb = g.get_partition_book()
if "trainer_id" in g.ndata:
train_nid = dgl.distributed.node_split(
g.ndata["train_mask"],
pb,
force_even=True,
node_trainer_ids=g.ndata["trainer_id"],
)
val_nid = dgl.distributed.node_split(
g.ndata["val_mask"],
pb,
force_even=True,
node_trainer_ids=g.ndata["trainer_id"],
)
test_nid = dgl.distributed.node_split(
g.ndata["test_mask"],
pb,
force_even=True,
node_trainer_ids=g.ndata["trainer_id"],
)
else:
train_nid = dgl.distributed.node_split(
g.ndata["train_mask"], pb, force_even=True
)
val_nid = dgl.distributed.node_split(
g.ndata["val_mask"], pb, force_even=True
)
test_nid = dgl.distributed.node_split(
g.ndata["test_mask"], pb, force_even=True
)
local_nid = pb.partid2nids(pb.partid).detach().numpy()
num_train_local = len(np.intersect1d(train_nid.numpy(), local_nid))
num_val_local = len(np.intersect1d(val_nid.numpy(), local_nid))
num_test_local = len(np.intersect1d(test_nid.numpy(), local_nid))
print(
f"part {g.rank()}, train: {len(train_nid)} (local: {num_train_local}), "
f"val: {len(val_nid)} (local: {num_val_local}), "
f"test: {len(test_nid)} (local: {num_test_local})"
)
del local_nid
if args.num_gpus == 0:
device = th.device("cpu")
else:
dev_id = g.rank() % args.num_gpus
device = th.device("cuda:" + str(dev_id))
n_classes = args.n_classes
if n_classes == 0:
labels = g.ndata["labels"][np.arange(g.num_nodes())]
n_classes = len(th.unique(labels[th.logical_not(th.isnan(labels))]))
del labels
print(f"Number of classes: {n_classes}")
# Pack data.
in_feats = g.ndata["features"].shape[1]
data = train_nid, val_nid, test_nid, in_feats, n_classes, g
# Train and evaluate.
epoch_time, test_acc = run(args, device, data)
print(
f"Summary of node classification(GraphSAGE): GraphName "
f"{args.graph_name} | TrainEpochTime(mean) {epoch_time:.4f} "
f"| TestAccuracy {test_acc:.4f}"
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Distributed GraphSAGE.")
parser.add_argument("--graph_name", type=str, help="graph name")
parser.add_argument(
"--ip_config", type=str, help="The file for IP configuration"
)
parser.add_argument(
"--part_config", type=str, help="The path to the partition config file"
)
parser.add_argument(
"--n_classes", type=int, default=0, help="the number of classes"
)
parser.add_argument(
"--backend",
type=str,
default="gloo",
help="pytorch distributed backend",
)
parser.add_argument(
"--num_gpus",
type=int,
default=0,
help="the number of GPU device. Use 0 for CPU training",
)
parser.add_argument("--num_epochs", type=int, default=20)
parser.add_argument("--num_hidden", type=int, default=16)
parser.add_argument("--num_layers", type=int, default=2)
parser.add_argument("--fan_out", type=str, default="10,25")
parser.add_argument("--batch_size", type=int, default=1000)
parser.add_argument("--batch_size_eval", type=int, default=100000)
parser.add_argument("--log_every", type=int, default=20)
parser.add_argument("--eval_every", type=int, default=5)
parser.add_argument("--lr", type=float, default=0.003)
parser.add_argument("--dropout", type=float, default=0.5)
parser.add_argument(
"--local_rank", type=int, help="get rank of the process"
)
parser.add_argument(
"--pad-data",
default=False,
action="store_true",
help="Pad train nid to the same length across machine, to ensure num "
"of batches to be the same.",
)
parser.add_argument(
"--use_graphbolt",
action="store_true",
help="Use GraphBolt for distributed train.",
)
args = parser.parse_args()
print(f"Arguments: {args}")
main(args)
examples/distributed/graphsage/node_classification_unsupervised.py
+476
View File
@@ -0,0 +1,476 @@
import argparse
import time
from contextlib import contextmanager
import dgl
import dgl.distributed
import dgl.function as fn
import dgl.nn.pytorch as dglnn
import numpy as np
import sklearn.linear_model as lm
import sklearn.metrics as skm
import torch as th
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import tqdm
class DistSAGE(nn.Module):
def __init__(
self, in_feats, n_hidden, n_classes, n_layers, activation, dropout
):
super().__init__()
self.n_layers = n_layers
self.n_hidden = n_hidden
self.n_classes = n_classes
self.layers = nn.ModuleList()
self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, "mean"))
for i in range(1, n_layers - 1):
self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, "mean"))
self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, "mean"))
self.dropout = nn.Dropout(dropout)
self.activation = activation
def forward(self, blocks, x):
h = x
for i, (layer, block) in enumerate(zip(self.layers, blocks)):
h = layer(block, h)
if i != len(self.layers) - 1:
h = self.activation(h)
h = self.dropout(h)
return h
def inference(self, g, x, batch_size, device):
"""
Inference with the GraphSAGE model on full neighbors (i.e. without
neighbor sampling).
g : the entire graph.
x : the input of entire node set.
The inference code is written in a fashion that it could handle any
number of nodes and layers.
"""
# During inference with sampling, multi-layer blocks are very
# inefficient because lots of computations in the first few layers are
# repeated. Therefore, we compute the representation of all nodes layer
# by layer. The nodes on each layer are of course splitted in batches.
# TODO: can we standardize this?
nodes = dgl.distributed.node_split(
np.arange(g.num_nodes()),
g.get_partition_book(),
force_even=True,
)
y = dgl.distributed.DistTensor(
(g.num_nodes(), self.n_hidden),
th.float32,
"h",
persistent=True,
)
for i, layer in enumerate(self.layers):
if i == len(self.layers) - 1:
y = dgl.distributed.DistTensor(
(g.num_nodes(), self.n_classes),
th.float32,
"h_last",
persistent=True,
)
# Create sampler
sampler = dgl.dataloading.NeighborSampler([-1])
# Create dataloader
dataloader = dgl.distributed.DistNodeDataLoader(
g,
nodes,
sampler,
batch_size=batch_size,
shuffle=False,
drop_last=False,
)
for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
block = blocks[0].to(device)
h = x[input_nodes].to(device)
h_dst = h[: block.number_of_dst_nodes()]
h = layer(block, (h, h_dst))
if i != len(self.layers) - 1:
h = self.activation(h)
h = self.dropout(h)
y[output_nodes] = h.cpu()
x = y
g.barrier()
return y
@contextmanager
def join(self):
"""dummy join for standalone"""
yield
def load_subtensor(g, input_nodes, device):
"""
Copys features and labels of a set of nodes onto GPU.
"""
batch_inputs = g.ndata["features"][input_nodes].to(device)
return batch_inputs
class CrossEntropyLoss(nn.Module):
def forward(self, block_outputs, pos_graph, neg_graph):
with pos_graph.local_scope():
pos_graph.ndata["h"] = block_outputs
pos_graph.apply_edges(fn.u_dot_v("h", "h", "score"))
pos_score = pos_graph.edata["score"]
with neg_graph.local_scope():
neg_graph.ndata["h"] = block_outputs
neg_graph.apply_edges(fn.u_dot_v("h", "h", "score"))
neg_score = neg_graph.edata["score"]
score = th.cat([pos_score, neg_score])
label = th.cat(
[th.ones_like(pos_score), th.zeros_like(neg_score)]
).long()
loss = F.binary_cross_entropy_with_logits(score, label.float())
return loss
def generate_emb(model, g, inputs, batch_size, device):
"""
Generate embeddings for each node
g : The entire graph.
inputs : The features of all the nodes.
batch_size : Number of nodes to compute at the same time.
device : The GPU device to evaluate on.
"""
model.eval()
with th.no_grad():
pred = model.inference(g, inputs, batch_size, device)
return pred
def compute_acc(emb, labels, train_nids, val_nids, test_nids):
"""
Compute the accuracy of prediction given the labels.
We will fist train a LogisticRegression model using the trained embeddings,
the training set, validation set and test set is provided as the arguments.
The final result is predicted by the lr model.
emb: The pretrained embeddings
labels: The ground truth
train_nids: The training set node ids
val_nids: The validation set node ids
test_nids: The test set node ids
"""
emb = emb[np.arange(labels.shape[0])].cpu().numpy()
train_nids = train_nids.cpu().numpy()
val_nids = val_nids.cpu().numpy()
test_nids = test_nids.cpu().numpy()
labels = labels.cpu().numpy()
emb = (emb - emb.mean(0, keepdims=True)) / emb.std(0, keepdims=True)
lr = lm.LogisticRegression(multi_class="multinomial", max_iter=10000)
lr.fit(emb[train_nids], labels[train_nids])
pred = lr.predict(emb)
eval_acc = skm.accuracy_score(labels[val_nids], pred[val_nids])
test_acc = skm.accuracy_score(labels[test_nids], pred[test_nids])
return eval_acc, test_acc
def run(args, device, data):
# Unpack data
(
train_eids,
train_nids,
in_feats,
g,
global_train_nid,
global_valid_nid,
global_test_nid,
labels,
) = data
# Create sampler
neg_sampler = dgl.dataloading.negative_sampler.Uniform(args.num_negs)
sampler = dgl.dataloading.NeighborSampler(
[int(fanout) for fanout in args.fan_out.split(",")]
)
# Create dataloader
exclude = "reverse_id" if args.remove_edge else None
reverse_eids = th.arange(g.num_edges()) if args.remove_edge else None
dataloader = dgl.distributed.DistEdgeDataLoader(
g,
train_eids,
sampler,
negative_sampler=neg_sampler,
exclude=exclude,
reverse_eids=reverse_eids,
batch_size=args.batch_size,
shuffle=True,
drop_last=False,
)
# Define model and optimizer
model = DistSAGE(
in_feats,
args.num_hidden,
args.num_hidden,
args.num_layers,
F.relu,
args.dropout,
)
model = model.to(device)
if not args.standalone:
if args.num_gpus == -1:
model = th.nn.parallel.DistributedDataParallel(model)
else:
dev_id = g.rank() % args.num_gpus
model = th.nn.parallel.DistributedDataParallel(
model, device_ids=[dev_id], output_device=dev_id
)
loss_fcn = CrossEntropyLoss()
loss_fcn = loss_fcn.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
# Training loop
epoch = 0
for epoch in range(args.num_epochs):
num_seeds = 0
num_inputs = 0
step_time = []
sample_t = []
feat_copy_t = []
forward_t = []
backward_t = []
update_t = []
iter_tput = []
start = time.time()
with model.join():
# Loop over the dataloader to sample the computation dependency
# graph as a list of blocks.
for step, (input_nodes, pos_graph, neg_graph, blocks) in enumerate(
dataloader
):
if args.debug:
# Verify exclude_edges functionality.
for block in blocks:
current_eids = block.edata[dgl.EID]
seed_eids = pos_graph.edata[dgl.EID]
if exclude is None:
assert th.any(th.isin(current_eids, seed_eids))
elif exclude == "self":
assert not th.any(th.isin(current_eids, seed_eids))
elif exclude == "reverse_id":
assert not th.any(th.isin(current_eids, seed_eids))
else:
raise ValueError(
f"Unsupported exclude type: {exclude}"
)
tic_step = time.time()
sample_t.append(tic_step - start)
copy_t = time.time()
pos_graph = pos_graph.to(device)
neg_graph = neg_graph.to(device)
blocks = [block.to(device) for block in blocks]
batch_inputs = load_subtensor(g, input_nodes, device)
copy_time = time.time()
feat_copy_t.append(copy_time - copy_t)
# Compute loss and prediction
batch_pred = model(blocks, batch_inputs)
loss = loss_fcn(batch_pred, pos_graph, neg_graph)
forward_end = time.time()
optimizer.zero_grad()
loss.backward()
compute_end = time.time()
forward_t.append(forward_end - copy_time)
backward_t.append(compute_end - forward_end)
# Aggregate gradients in multiple nodes.
optimizer.step()
update_t.append(time.time() - compute_end)
pos_edges = pos_graph.num_edges()
step_t = time.time() - start
step_time.append(step_t)
iter_tput.append(pos_edges / step_t)
num_seeds += pos_edges
if step % args.log_every == 0:
print(
"[{}] Epoch {:05d} | Step {:05d} | Loss {:.4f} | Speed "
"(samples/sec) {:.4f} | time {:.3f}s | sample {:.3f} | "
"copy {:.3f} | forward {:.3f} | backward {:.3f} | "
"update {:.3f}".format(
g.rank(),
epoch,
step,
loss.item(),
np.mean(iter_tput[3:]),
np.sum(step_time[-args.log_every :]),
np.sum(sample_t[-args.log_every :]),
np.sum(feat_copy_t[-args.log_every :]),
np.sum(forward_t[-args.log_every :]),
np.sum(backward_t[-args.log_every :]),
np.sum(update_t[-args.log_every :]),
)
)
start = time.time()
print(
"[{}]Epoch Time(s): {:.4f}, sample: {:.4f}, data copy: {:.4f}, "
"forward: {:.4f}, backward: {:.4f}, update: {:.4f}, #seeds: {}, "
"#inputs: {}".format(
g.rank(),
np.sum(step_time),
np.sum(sample_t),
np.sum(feat_copy_t),
np.sum(forward_t),
np.sum(backward_t),
np.sum(update_t),
num_seeds,
num_inputs,
)
)
epoch += 1
# evaluate the embedding using LogisticRegression
pred = generate_emb(
model if args.standalone else model.module,
g,
g.ndata["features"],
args.batch_size_eval,
device,
)
if g.rank() == 0:
eval_acc, test_acc = compute_acc(
pred, labels, global_train_nid, global_valid_nid, global_test_nid
)
print("eval acc {:.4f}; test acc {:.4f}".format(eval_acc, test_acc))
# sync for eval and test
if not args.standalone:
th.distributed.barrier()
if not args.standalone:
g._client.barrier()
# save features into file
if g.rank() == 0:
th.save(pred, "emb.pt")
else:
th.save(pred, "emb.pt")
def main(args):
print("--- Distributed node classification with GraphSAGE unsuperised ---")
dgl.distributed.initialize(args.ip_config)
if not args.standalone:
th.distributed.init_process_group(backend="gloo")
g = dgl.distributed.DistGraph(args.graph_name, part_config=args.part_config)
print("rank:", g.rank())
print("number of edges", g.num_edges())
train_eids = dgl.distributed.edge_split(
th.ones((g.num_edges(),), dtype=th.bool),
g.get_partition_book(),
force_even=True,
)
train_nids = dgl.distributed.node_split(
th.ones((g.num_nodes(),), dtype=th.bool), g.get_partition_book()
)
global_train_nid = th.LongTensor(
np.nonzero(g.ndata["train_mask"][np.arange(g.num_nodes())])
)
global_valid_nid = th.LongTensor(
np.nonzero(g.ndata["val_mask"][np.arange(g.num_nodes())])
)
global_test_nid = th.LongTensor(
np.nonzero(g.ndata["test_mask"][np.arange(g.num_nodes())])
)
labels = g.ndata["labels"][np.arange(g.num_nodes())]
if args.num_gpus == -1:
device = th.device("cpu")
else:
dev_id = g.rank() % args.num_gpus
device = th.device("cuda:" + str(dev_id))
# Pack data
in_feats = g.ndata["features"].shape[1]
global_train_nid = global_train_nid.squeeze()
global_valid_nid = global_valid_nid.squeeze()
global_test_nid = global_test_nid.squeeze()
print("number of train {}".format(global_train_nid.shape[0]))
print("number of valid {}".format(global_valid_nid.shape[0]))
print("number of test {}".format(global_test_nid.shape[0]))
data = (
train_eids,
train_nids,
in_feats,
g,
global_train_nid,
global_valid_nid,
global_test_nid,
labels,
)
run(args, device, data)
print("parent ends")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="GCN")
parser.add_argument("--graph_name", type=str, help="graph name")
parser.add_argument("--id", type=int, help="the partition id")
parser.add_argument(
"--ip_config", type=str, help="The file for IP configuration"
)
parser.add_argument(
"--part_config", type=str, help="The path to the partition config file"
)
parser.add_argument("--n_classes", type=int, help="the number of classes")
parser.add_argument(
"--num_gpus",
type=int,
default=-1,
help="the number of GPU device. Use -1 for CPU training",
)
parser.add_argument("--num_epochs", type=int, default=20)
parser.add_argument("--num_hidden", type=int, default=16)
parser.add_argument("--num-layers", type=int, default=2)
parser.add_argument("--fan_out", type=str, default="10,25")
parser.add_argument("--batch_size", type=int, default=1000)
parser.add_argument("--batch_size_eval", type=int, default=100000)
parser.add_argument("--log_every", type=int, default=20)
parser.add_argument("--eval_every", type=int, default=5)
parser.add_argument("--lr", type=float, default=0.003)
parser.add_argument("--dropout", type=float, default=0.5)
parser.add_argument(
"--local_rank", type=int, help="get rank of the process"
)
parser.add_argument(
"--standalone", action="store_true", help="run in the standalone mode"
)
parser.add_argument("--num_negs", type=int, default=1)
parser.add_argument(
"--remove_edge",
default=False,
action="store_true",
help="whether to remove edges during sampling",
)
parser.add_argument(
"--debug",
default=False,
action="store_true",
help="whether to verify functionality of remove edges",
)
args = parser.parse_args()
print(args)
main(args)
examples/distributed/graphsage/partition_graph.py
+136
View File
@@ -0,0 +1,136 @@
import argparse
import time
import dgl
import torch as th
from dgl.data import RedditDataset
from ogb.nodeproppred import DglNodePropPredDataset
def load_reddit(self_loop=True):
"""Load reddit dataset."""
data = RedditDataset(self_loop=self_loop)
g = data[0]
g.ndata["features"] = g.ndata.pop("feat")
g.ndata["labels"] = g.ndata.pop("label")
return g, data.num_classes
def load_ogb(name, root="dataset"):
"""Load ogbn dataset."""
data = DglNodePropPredDataset(name=name, root=root)
splitted_idx = data.get_idx_split()
graph, labels = data[0]
labels = labels[:, 0]
graph.ndata["features"] = graph.ndata.pop("feat")
graph.ndata["labels"] = labels
num_labels = len(th.unique(labels[th.logical_not(th.isnan(labels))]))
# Find the node IDs in the training, validation, and test set.
train_nid, val_nid, test_nid = (
splitted_idx["train"],
splitted_idx["valid"],
splitted_idx["test"],
)
train_mask = th.zeros((graph.num_nodes(),), dtype=th.bool)
train_mask[train_nid] = True
val_mask = th.zeros((graph.num_nodes(),), dtype=th.bool)
val_mask[val_nid] = True
test_mask = th.zeros((graph.num_nodes(),), dtype=th.bool)
test_mask[test_nid] = True
graph.ndata["train_mask"] = train_mask
graph.ndata["val_mask"] = val_mask
graph.ndata["test_mask"] = test_mask
return graph, num_labels
if __name__ == "__main__":
argparser = argparse.ArgumentParser("Partition graph")
argparser.add_argument(
"--dataset",
type=str,
default="reddit",
help="datasets: reddit, ogbn-products, ogbn-papers100M",
)
argparser.add_argument(
"--num_parts", type=int, default=4, help="number of partitions"
)
argparser.add_argument(
"--part_method", type=str, default="metis", help="the partition method"
)
argparser.add_argument(
"--balance_train",
action="store_true",
help="balance the training size in each partition.",
)
argparser.add_argument(
"--undirected",
action="store_true",
help="turn the graph into an undirected graph.",
)
argparser.add_argument(
"--balance_edges",
action="store_true",
help="balance the number of edges in each partition.",
)
argparser.add_argument(
"--num_trainers_per_machine",
type=int,
default=1,
help="the number of trainers per machine. The trainer ids are stored\
in the node feature 'trainer_id'",
)
argparser.add_argument(
"--output",
type=str,
default="data",
help="Output path of partitioned graph.",
)
argparser.add_argument(
"--use_graphbolt",
action="store_true",
help="Use GraphBolt for distributed train.",
)
args = argparser.parse_args()
start = time.time()
if args.dataset == "reddit":
g, _ = load_reddit()
elif args.dataset in ["ogbn-products", "ogbn-papers100M"]:
g, _ = load_ogb(args.dataset)
else:
raise RuntimeError(f"Unknown dataset: {args.dataset}")
print(
"Load {} takes {:.3f} seconds".format(args.dataset, time.time() - start)
)
print("|V|={}, |E|={}".format(g.num_nodes(), g.num_edges()))
print(
"train: {}, valid: {}, test: {}".format(
th.sum(g.ndata["train_mask"]),
th.sum(g.ndata["val_mask"]),
th.sum(g.ndata["test_mask"]),
)
)
if args.balance_train:
balance_ntypes = g.ndata["train_mask"]
else:
balance_ntypes = None
if args.undirected:
sym_g = dgl.to_bidirected(g, readonly=True)
for key in g.ndata:
sym_g.ndata[key] = g.ndata[key]
g = sym_g
dgl.distributed.partition_graph(
g,
args.dataset,
args.num_parts,
args.output,
part_method=args.part_method,
balance_ntypes=balance_ntypes,
balance_edges=args.balance_edges,
num_trainers_per_machine=args.num_trainers_per_machine,
use_graphbolt=args.use_graphbolt,
)
examples/distributed/rgcn/README.md
+247
View File
@@ -0,0 +1,247 @@
## Distributed training
This is an example of training RGCN node classification in a distributed fashion. Currently, the example train RGCN graphs with input node features.
Before training, install python libs by pip:
```bash
pip3 install ogb pyarrow
```
To train RGCN, it has four steps:
### Step 0: Setup a Distributed File System
* You may skip this step if your cluster already has folder(s) synchronized across machines.
To perform distributed training, files and codes need to be accessed across multiple machines. A distributed file system would perfectly handle the job (i.e., NFS, Ceph).
#### Server side setup
Here is an example of how to setup NFS. First, install essential libs on the storage server
```bash
sudo apt-get install nfs-kernel-server
```
Below we assume the user account is `ubuntu` and we create a directory of `workspace` in the home directory.
```bash
mkdir -p /home/ubuntu/workspace
```
We assume that the all servers are under a subnet with ip range `192.168.0.0` to `192.168.255.255`. The exports configuration needs to be modifed to
```bash
sudo vim /etc/exports
# add the following line
/home/ubuntu/workspace 192.168.0.0/16(rw,sync,no_subtree_check)
```
The server's internal ip can be checked via `ifconfig` or `ip`. If the ip does not begin with `192.168`, then you may use
```bash
# for ip range 10.0.0.0 - 10.255.255.255
/home/ubuntu/workspace 10.0.0.0/8(rw,sync,no_subtree_check)
# for ip range 172.16.0.0 - 172.31.255.255
/home/ubuntu/workspace 172.16.0.0/12(rw,sync,no_subtree_check)
```
Then restart NFS, the setup on server side is finished.
```
sudo systemctl restart nfs-kernel-server
```
For configraution details, please refer to [NFS ArchWiki](https://wiki.archlinux.org/index.php/NFS).
#### Client side setup
To use NFS, clients also require to install essential packages
```
sudo apt-get install nfs-common
```
You can either mount the NFS manually
```
mkdir -p /home/ubuntu/workspace
sudo mount -t nfs <nfs-server-ip>:/home/ubuntu/workspace /home/ubuntu/workspace
```
or edit the fstab so the folder will be mounted automatically
```
# vim /etc/fstab
## append the following line to the file
<nfs-server-ip>:/home/ubuntu/workspace /home/ubuntu/workspace nfs defaults 0 0
```
Then run `mount -a`.
Now go to `/home/ubuntu/workspace` and clone the DGL Github repository.
### Step 1: set IP configuration file.
User need to set their own IP configuration file `ip_config.txt` before training. For example, if we have four machines in current cluster, the IP configuration could like this:
```bash
172.31.0.1
172.31.0.2
```
Users need to make sure that the master node (node-0) has right permission to ssh to all the other nodes without password authentication.
[This link](https://linuxize.com/post/how-to-setup-passwordless-ssh-login/) provides instructions of setting passwordless SSH login.
### Step 2: partition the graph.
The example provides a script to partition some builtin graphs such as ogbn-mag graph.
If we want to train RGCN on 2 machines, we need to partition the graph into 2 parts.
In this example, we partition the ogbn-mag graph into 2 parts with Metis. The partitions are balanced with respect to the number of nodes, the number of edges and the number of labelled nodes.
```bash
python3 partition_graph.py --dataset ogbn-mag --num_parts 2 --balance_train --balance_edges
```
If we want to train RGCN with `GraphBolt`, we need to append `--use_graphbolt` to generate partitions in `GraphBolt` format.
```bash
python3 partition_graph.py --dataset ogbn-mag --num_parts 2 --balance_train --balance_edges --use_graphbolt
```
If we have already partitioned into `DGL` format, just convert them directly like below:
```
python3 -c "import dgl; dgl.distributed.dgl_partition_to_graphbolt('ogbn-products.json')"
```
### Step 3: Launch distributed jobs
DGL provides a script to launch the training job in the cluster. `part_config` and `ip_config`
specify relative paths to the path of the workspace.
The command below launches 4 training processes on each machine as we'd like to utilize 4 GPUs for training.
```bash
python3 ~/workspace/dgl/tools/launch.py \
--workspace ~/workspace/dgl/examples/distributed/rgcn/ \
--num_trainers 4 \
--num_servers 2 \
--num_samplers 0 \
--part_config data/ogbn-mag.json \
--ip_config ip_config.txt \
"python3 node_classification.py --graph-name ogbn-mag --dataset ogbn-mag --fanout='25,25' --batch-size 1024 --n-hidden 64 --lr 0.01 --eval-batch-size 1024 --low-mem --dropout 0.5 --use-self-loop --n-bases 2 --n-epochs 3 --layer-norm --ip-config ip_config.txt --num_gpus 4"
```
If we want to train RGCN with `GraphBolt`, we need to append `--use_graphbolt`.
```bash
python3 ~/workspace/dgl/tools/launch.py \
--workspace ~/workspace/dgl/examples/distributed/rgcn/ \
--num_trainers 4 \
--num_servers 2 \
--num_samplers 0 \
--part_config data/ogbn-mag.json \
--ip_config ip_config.txt \
"python3 node_classification.py --graph-name ogbn-mag --dataset ogbn-mag --fanout='25,25' --batch-size 1024 --n-hidden 64 --lr 0.01 --eval-batch-size 1024 --low-mem --dropout 0.5 --use-self-loop --n-bases 2 --n-epochs 3 --layer-norm --ip-config ip_config.txt --num_gpus 4 --use_graphbolt"
```
**Note:** if you are using conda or other virtual environments on the remote machines, you need to replace `python3` in the command string (i.e. the last argument) with the path to the Python interpreter in that environment.
## Comparison between `DGL` and `GraphBolt`
### Partition sizes
Compared to `DGL`, `GraphBolt` partitions are reduced to **19%** for `ogbn-mag`.
`ogbn-mag`
| Data Formats | File Name | Part 0 | Part 1 |
| ------------ | ---------------------------- | ------ | ------ |
| DGL | graph.dgl | 714MB | 716MB |
| GraphBolt | fused_csc_sampling_graph.pt | 137MB | 136MB |
### Performance
Compared to `DGL`, `GraphBolt`'s sampler works faster(reduced to **16%** `ogbn-mag`). `Min` and `Max` are statistics of all trainers on all nodes(machines).
As for RAM usage, the shared memory(measured by **shared** field of `free` command) usage decreases due to smaller graph partitions in `GraphBolt`. The peak memory used by processes(measured by **used** field of `free` command) decreases as well.
`ogbn-mag`
| Data Formats | Sample Time Per Epoch (CPU) | Test Accuracy (3 epochs) | shared | used (peak) | CPU Util |
| ------------ | --------------------------- | ------------------------- | ----- | ---- | ----- |
| DGL | Min: 48.2s, Max: 91.4s | 42.76% | 1.3GB | 9.2GB| 10.4% |
| GraphBolt | Min: 9.2s, Max: 11.9s | 42.46% | 742MB | 5.9GB| 18.1% |
## Demonstrate and profile sampling for Link Prediction task
### DGL
```
python3 ~/workspace/dgl/tools/launch.py \
--workspace ~/workspace/dgl/examples/distributed/rgcn/ \
--num_trainers 4 \
--num_servers 2 \
--num_samplers 0 \
--part_config ~/data/ogbn_mag_lp/ogbn-mag.json \
--ip_config ~/workspace/ip_config.txt \
"python3 lp_perf.py --fanout='25,25' --batch-size 1024 --n-epochs 1 --graph-name ogbn-mag --ip-config ~/workspace/ip_config.txt --num_gpus 4 --remove_edge"
```
### GraphBolt
In order to sample with `GraphBolt`, we need to convert partitions into `GraphBolt` formats with below command.
```
python3 -c "import dgl;dgl.distributed.dgl_partition_to_graphbolt('/home/ubuntu/workspace/data/ogbn_mag_lp/ogbn-mag.json', store_eids=True, graph_formats='coo')"
```
Then train with appended `--use_graphbolt`.
```
python3 ~/workspace/dgl/tools/launch.py \
--workspace ~/workspace/dgl/examples/distributed/rgcn/ \
--num_trainers 4 \
--num_servers 2 \
--num_samplers 0 \
--part_config ~/data/ogbn_mag_lp/ogbn-mag.json \
--ip_config ~/workspace/ip_config.txt \
"python3 lp_perf.py --fanout='25,25' --batch-size 1024 --n-epochs 1 --graph-name ogbn-mag --ip-config ~/workspace/ip_config.txt --num_gpus 4 --remove_edge --use_graphbolt"
```
### Partition sizes
Compared to `DGL`, `GraphBolt` partitions are reduced to **72%** for `ogbn-mag`.
#### ogbn-mag
| Data Formats | File Name | Part 0 | Part 1 |
| ------------ | ---------------------------- | ------ | ------ |
| DGL | graph.dgl | 714MB | 716MB |
| GraphBolt | fused_csc_sampling_graph.pt | 512MB | 514MB |
### Performance Comparison
#### Major used parameters
1. 2 nodes(g4dn.metal), 4 trainers, 2 servers per node. Sample on main process.
2. 2 layers.
3. fanouts = 25, 25 for all edge types.
4. batch_size = 1024.
5. seed edge IDs are all edges of ("author", "writes", "paper"), ~7M in total.
6. ratio of negative sampler = 3.
7. exclude = "reverse_types".
#### ogbn-mag
Compared to `DGL`, sampling with `GraphBolt` is reduced to **15%**. As for the overhead of `exclude`, it's about **5%** in this test. This number could be higher if larger `fanout` or `batch size` is applied.
The time shown below is the mean sampling time per iteration(60 iters in total, slowest rank). Unit: seconds
| Data Formats | No Exclude | Exclude |
| ------------ | ---------- | ------- |
| DGL | 6.50 | 6.86 |
| GraphBolt | 0.95 | 1.00 |
examples/distributed/rgcn/lp_perf.py
+252
View File
@@ -0,0 +1,252 @@
"""
[For internal use only]
Demonstrate and profile the performance of sampling for link prediction tasks.
"""
import argparse
import time
import dgl
import numpy as np
import torch as th
def run(args, g, train_eids):
fanouts = [int(fanout) for fanout in args.fanout.split(",")]
neg_sampler = dgl.dataloading.negative_sampler.Uniform(3)
prob = args.prob_or_mask
sampler = dgl.dataloading.MultiLayerNeighborSampler(
fanouts,
prob=prob,
)
exclude = None
reverse_etypes = None
if args.remove_edge:
exclude = "reverse_types"
# add reverse edge types mapping.
reverse_etypes = {
("author", "affiliated_with", "institution"): (
"institution",
"rev-affiliated_with",
"author",
),
("author", "writes", "paper"): ("paper", "rev-writes", "author"),
("paper", "has_topic", "field_of_study"): (
"field_of_study",
"rev-has_topic",
"paper",
),
("paper", "cites", "paper"): ("paper", "rev-cites", "paper"),
("institution", "rev-affiliated_with", "author"): (
"author",
"affiliated_with",
"institution",
),
("paper", "rev-writes", "author"): ("author", "writes", "paper"),
("field_of_study", "rev-has_topic", "paper"): (
"paper",
"has_topic",
"field_of_study",
),
("paper", "rev-cites", "paper"): ("paper", "cites", "paper"),
}
dataloader = dgl.dataloading.DistEdgeDataLoader(
g,
train_eids,
sampler,
negative_sampler=neg_sampler,
exclude=exclude,
reverse_etypes=reverse_etypes,
batch_size=args.batch_size,
shuffle=True,
drop_last=False,
)
for epoch in range(args.n_epochs):
sample_times = []
tic = time.time()
epoch_tic = time.time()
for step, sample_data in enumerate(dataloader):
input_nodes, pos_graph, neg_graph, blocks = sample_data
if args.debug:
# Verify prob/mask values.
for block in blocks:
for c_etype in block.canonical_etypes:
homo_eids = block.edges[c_etype].data[dgl.EID]
assert th.all(
g.edges[c_etype].data[prob][homo_eids] > 0
)
# Verify exclude_edges functionality.
current_eids = blocks[-1].edata[dgl.EID]
seed_eids = pos_graph.edata[dgl.EID]
if exclude is None:
assert th.any(th.isin(current_eids, seed_eids))
elif exclude == "self":
assert not th.any(th.isin(current_eids, seed_eids))
elif exclude == "reverse_id":
assert not th.any(th.isin(current_eids, seed_eids))
elif exclude == "reverse_types":
for src_type, etype, dst_type in pos_graph.canonical_etypes:
reverse_etype = reverse_etypes[
(src_type, etype, dst_type)
]
seed_eids = pos_graph.edges[etype].data[dgl.EID]
if (src_type, etype, dst_type) in blocks[
-1
].canonical_etypes:
assert not th.any(
th.isin(
blocks[-1].edges[etype].data[dgl.EID],
seed_eids,
)
)
if reverse_etype in blocks[-1].canonical_etypes:
assert not th.any(
th.isin(
blocks[-1]
.edges[reverse_etype]
.data[dgl.EID],
seed_eids,
)
)
else:
raise ValueError(f"Unsupported exclude type: {exclude}")
sample_times.append(time.time() - tic)
if step % 10 == 0:
print(
f"[{g.rank()}]Epoch {epoch} | Step {step} | Sample Time {np.mean(sample_times[10:]):.4f}"
)
tic = time.time()
print(
f"[{g.rank()}]Epoch {epoch} | Total time {time.time() - epoch_tic} | Sample Time {np.mean(sample_times[100:]):.4f}"
)
g.barrier()
def rand_init_prob(shape, dtype):
prob = th.rand(shape)
prob[th.randperm(len(prob))[: int(len(prob) * 0.5)]] = 0.0
return prob
def rand_init_mask(shape, dtype):
prob = th.rand(shape)
prob[th.randperm(len(prob))[: int(len(prob) * 0.5)]] = 0.0
return (prob > 0.2).to(th.float32)
def main(args):
dgl.distributed.initialize(args.ip_config, use_graphbolt=args.use_graphbolt)
backend = "gloo" if args.num_gpus == -1 else "nccl"
th.distributed.init_process_group(backend=backend)
g = dgl.distributed.DistGraph(args.graph_name)
print("rank:", g.rank())
# Assign prob/masks to edges.
for c_etype in g.canonical_etypes:
shape = (g.num_edges(etype=c_etype),)
g.edges[c_etype].data["prob"] = dgl.distributed.DistTensor(
shape,
th.float32,
init_func=rand_init_prob,
part_policy=g.get_edge_partition_policy(c_etype),
)
g.edges[c_etype].data["mask"] = dgl.distributed.DistTensor(
shape,
th.float32,
init_func=rand_init_mask,
part_policy=g.get_edge_partition_policy(c_etype),
)
pb = g.get_partition_book()
c_etype = ("author", "writes", "paper")
train_eids = dgl.distributed.edge_split(
th.ones((g.num_edges(etype=c_etype),), dtype=th.bool),
g.get_partition_book(),
etype=c_etype,
force_even=True,
)
train_eids = {c_etype: train_eids}
local_eids = pb.partid2eids(pb.partid, c_etype).detach().numpy()
print(
"part {}, train: {} (local: {})".format(
g.rank(),
len(train_eids[c_etype]),
len(np.intersect1d(train_eids[c_etype].numpy(), local_eids)),
)
)
run(
args,
g,
train_eids,
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Sampling Performance Profiling For Link Prediction Tasks"
)
parser.add_argument("--graph-name", type=str, help="graph name")
parser.add_argument(
"--ip-config", type=str, help="The file for IP configuration"
)
parser.add_argument(
"--num_gpus",
type=int,
default=-1,
help="the number of GPU device. Use -1 for CPU training",
)
parser.add_argument(
"-e",
"--n-epochs",
type=int,
default=5,
help="number of training epochs",
)
parser.add_argument(
"--fanout",
type=str,
default="4, 4",
help="Fan-out of neighbor sampling.",
)
parser.add_argument(
"--batch-size", type=int, default=100, help="Mini-batch size. "
)
parser.add_argument(
"--use_graphbolt",
default=False,
action="store_true",
help="Use GraphBolt for distributed train.",
)
parser.add_argument(
"--remove_edge",
default=False,
action="store_true",
help="whether to remove edges during sampling",
)
parser.add_argument(
"--debug",
default=False,
action="store_true",
help="whether to remove edges during sampling",
)
parser.add_argument(
"--prob_or_mask",
type=str,
default="prob",
help="whether to use prob or mask during sampling",
)
args = parser.parse_args()
print(args)
main(args)
examples/distributed/rgcn/node_classification.py
+927
View File
@@ -0,0 +1,927 @@
"""
Modeling Relational Data with Graph Convolutional Networks
Paper: https://arxiv.org/abs/1703.06103
Code: https://github.com/tkipf/relational-gcn
Difference compared to tkipf/relation-gcn
* l2norm applied to all weights
* remove nodes that won't be touched
"""
import argparse
import gc, os
import itertools
import time
import numpy as np
os.environ["DGLBACKEND"] = "pytorch"
from functools import partial
import dgl
import dgl.distributed
import torch as th
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F
import tqdm
from dgl import DGLGraph, nn as dglnn
from dgl.distributed import DistDataLoader
from ogb.nodeproppred import DglNodePropPredDataset
from torch.multiprocessing import Queue
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader
class RelGraphConvLayer(nn.Module):
r"""Relational graph convolution layer.
Parameters
----------
in_feat : int
Input feature size.
out_feat : int
Output feature size.
rel_names : list[str]
Relation names.
num_bases : int, optional
Number of bases. If is none, use number of relations. Default: None.
weight : bool, optional
True if a linear layer is applied after message passing. Default: True
bias : bool, optional
True if bias is added. Default: True
activation : callable, optional
Activation function. Default: None
self_loop : bool, optional
True to include self loop message. Default: False
dropout : float, optional
Dropout rate. Default: 0.0
"""
def __init__(
self,
in_feat,
out_feat,
rel_names,
num_bases,
*,
weight=True,
bias=True,
activation=None,
self_loop=False,
dropout=0.0
):
super(RelGraphConvLayer, self).__init__()
self.in_feat = in_feat
self.out_feat = out_feat
self.rel_names = rel_names
self.num_bases = num_bases
self.bias = bias
self.activation = activation
self.self_loop = self_loop
self.conv = dglnn.HeteroGraphConv(
{
rel: dglnn.GraphConv(
in_feat, out_feat, norm="right", weight=False, bias=False
)
for rel in rel_names
}
)
self.use_weight = weight
self.use_basis = num_bases < len(self.rel_names) and weight
if self.use_weight:
if self.use_basis:
self.basis = dglnn.WeightBasis(
(in_feat, out_feat), num_bases, len(self.rel_names)
)
else:
self.weight = nn.Parameter(
th.Tensor(len(self.rel_names), in_feat, out_feat)
)
nn.init.xavier_uniform_(
self.weight, gain=nn.init.calculate_gain("relu")
)
# bias
if bias:
self.h_bias = nn.Parameter(th.Tensor(out_feat))
nn.init.zeros_(self.h_bias)
# weight for self loop
if self.self_loop:
self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
nn.init.xavier_uniform_(
self.loop_weight, gain=nn.init.calculate_gain("relu")
)
self.dropout = nn.Dropout(dropout)
def forward(self, g, inputs):
"""Forward computation
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.
"""
g = g.local_var()
if self.use_weight:
weight = self.basis() if self.use_basis else self.weight
wdict = {
self.rel_names[i]: {"weight": w.squeeze(0)}
for i, w in enumerate(th.split(weight, 1, dim=0))
}
else:
wdict = {}
if g.is_block:
inputs_src = inputs
inputs_dst = {
k: v[: g.number_of_dst_nodes(k)] for k, v in inputs.items()
}
else:
inputs_src = inputs_dst = inputs
hs = self.conv(g, inputs, mod_kwargs=wdict)
def _apply(ntype, h):
if self.self_loop:
h = h + th.matmul(inputs_dst[ntype], self.loop_weight)
if self.bias:
h = h + self.h_bias
if self.activation:
h = self.activation(h)
return self.dropout(h)
return {ntype: _apply(ntype, h) for ntype, h in hs.items()}
class EntityClassify(nn.Module):
"""Entity classification class for RGCN
Parameters
----------
device : int
Device to run the layer.
num_nodes : int
Number of nodes.
h_dim : int
Hidden dim size.
out_dim : int
Output dim size.
rel_names : list of str
A list of relation names.
num_bases : int
Number of bases. If is none, use number of relations.
num_hidden_layers : int
Number of hidden RelGraphConv Layer
dropout : float
Dropout
use_self_loop : bool
Use self loop if True, default False.
"""
def __init__(
self,
device,
h_dim,
out_dim,
rel_names,
num_bases=None,
num_hidden_layers=1,
dropout=0,
use_self_loop=False,
layer_norm=False,
):
super(EntityClassify, self).__init__()
self.device = device
self.h_dim = h_dim
self.out_dim = out_dim
self.num_bases = None if num_bases < 0 else num_bases
self.num_hidden_layers = num_hidden_layers
self.dropout = dropout
self.use_self_loop = use_self_loop
self.layer_norm = layer_norm
self.layers = nn.ModuleList()
# i2h
self.layers.append(
RelGraphConvLayer(
self.h_dim,
self.h_dim,
rel_names,
self.num_bases,
activation=F.relu,
self_loop=self.use_self_loop,
dropout=self.dropout,
)
)
# h2h
for idx in range(self.num_hidden_layers):
self.layers.append(
RelGraphConvLayer(
self.h_dim,
self.h_dim,
rel_names,
self.num_bases,
activation=F.relu,
self_loop=self.use_self_loop,
dropout=self.dropout,
)
)
# h2o
self.layers.append(
RelGraphConvLayer(
self.h_dim,
self.out_dim,
rel_names,
self.num_bases,
activation=None,
self_loop=self.use_self_loop,
)
)
def forward(self, blocks, feats, norm=None):
if blocks is None:
# full graph training
blocks = [self.g] * len(self.layers)
h = feats
for layer, block in zip(self.layers, blocks):
block = block.to(self.device)
h = layer(block, h)
return h
def init_emb(shape, dtype):
arr = th.zeros(shape, dtype=dtype)
nn.init.uniform_(arr, -1.0, 1.0)
return arr
class DistEmbedLayer(nn.Module):
r"""Embedding layer for featureless heterograph.
Parameters
----------
dev_id : int
Device to run the layer.
g : DistGraph
training graph
embed_size : int
Output embed size
sparse_emb: bool
Whether to use sparse embedding
Default: False
dgl_sparse_emb: bool
Whether to use DGL sparse embedding
Default: False
embed_name : str, optional
Embed name
"""
def __init__(
self,
dev_id,
g,
embed_size,
sparse_emb=False,
dgl_sparse_emb=False,
feat_name="feat",
embed_name="node_emb",
):
super(DistEmbedLayer, self).__init__()
self.dev_id = dev_id
self.embed_size = embed_size
self.embed_name = embed_name
self.feat_name = feat_name
self.sparse_emb = sparse_emb
self.g = g
self.ntype_id_map = {g.get_ntype_id(ntype): ntype for ntype in g.ntypes}
self.node_projs = nn.ModuleDict()
for ntype in g.ntypes:
if feat_name in g.nodes[ntype].data:
self.node_projs[ntype] = nn.Linear(
g.nodes[ntype].data[feat_name].shape[1], embed_size
)
nn.init.xavier_uniform_(self.node_projs[ntype].weight)
print("node {} has data {}".format(ntype, feat_name))
if sparse_emb:
if dgl_sparse_emb:
self.node_embeds = {}
for ntype in g.ntypes:
# We only create embeddings for nodes without node features.
if feat_name not in g.nodes[ntype].data:
part_policy = g.get_node_partition_policy(ntype)
self.node_embeds[ntype] = dgl.distributed.DistEmbedding(
g.num_nodes(ntype),
self.embed_size,
embed_name + "_" + ntype,
init_emb,
part_policy,
)
else:
self.node_embeds = nn.ModuleDict()
for ntype in g.ntypes:
# We only create embeddings for nodes without node features.
if feat_name not in g.nodes[ntype].data:
self.node_embeds[ntype] = th.nn.Embedding(
g.num_nodes(ntype),
self.embed_size,
sparse=self.sparse_emb,
)
nn.init.uniform_(
self.node_embeds[ntype].weight, -1.0, 1.0
)
else:
self.node_embeds = nn.ModuleDict()
for ntype in g.ntypes:
# We only create embeddings for nodes without node features.
if feat_name not in g.nodes[ntype].data:
self.node_embeds[ntype] = th.nn.Embedding(
g.num_nodes(ntype), self.embed_size
)
nn.init.uniform_(self.node_embeds[ntype].weight, -1.0, 1.0)
def forward(self, node_ids):
"""Forward computation
Parameters
----------
node_ids : dict of Tensor
node ids to generate embedding for.
Returns
-------
tensor
embeddings as the input of the next layer
"""
embeds = {}
for ntype in node_ids:
if self.feat_name in self.g.nodes[ntype].data:
embeds[ntype] = self.node_projs[ntype](
self.g.nodes[ntype]
.data[self.feat_name][node_ids[ntype]]
.to(self.dev_id)
)
else:
embeds[ntype] = self.node_embeds[ntype](node_ids[ntype]).to(
self.dev_id
)
return embeds
def compute_acc(results, labels):
"""
Compute the accuracy of prediction given the labels.
"""
labels = labels.long()
return (results == labels).float().sum() / len(results)
def evaluate(
g,
model,
embed_layer,
labels,
eval_loader,
test_loader,
all_val_nid,
all_test_nid,
):
model.eval()
embed_layer.eval()
eval_logits = []
eval_seeds = []
global_results = dgl.distributed.DistTensor(
labels.shape, th.long, "results", persistent=True
)
with th.no_grad():
th.cuda.empty_cache()
for sample_data in tqdm.tqdm(eval_loader):
input_nodes, seeds, blocks = sample_data
seeds = seeds["paper"]
feats = embed_layer(input_nodes)
logits = model(blocks, feats)
assert len(logits) == 1
logits = logits["paper"]
eval_logits.append(logits.cpu().detach())
assert np.all(seeds.numpy() < g.num_nodes("paper"))
eval_seeds.append(seeds.cpu().detach())
eval_logits = th.cat(eval_logits)
eval_seeds = th.cat(eval_seeds)
global_results[eval_seeds] = eval_logits.argmax(dim=1)
test_logits = []
test_seeds = []
with th.no_grad():
th.cuda.empty_cache()
for sample_data in tqdm.tqdm(test_loader):
input_nodes, seeds, blocks = sample_data
seeds = seeds["paper"]
feats = embed_layer(input_nodes)
logits = model(blocks, feats)
assert len(logits) == 1
logits = logits["paper"]
test_logits.append(logits.cpu().detach())
assert np.all(seeds.numpy() < g.num_nodes("paper"))
test_seeds.append(seeds.cpu().detach())
test_logits = th.cat(test_logits)
test_seeds = th.cat(test_seeds)
global_results[test_seeds] = test_logits.argmax(dim=1)
g.barrier()
if g.rank() == 0:
return compute_acc(
global_results[all_val_nid], labels[all_val_nid]
), compute_acc(global_results[all_test_nid], labels[all_test_nid])
else:
return -1, -1
def run(args, device, data):
(
g,
num_classes,
train_nid,
val_nid,
test_nid,
labels,
all_val_nid,
all_test_nid,
) = data
fanouts = [int(fanout) for fanout in args.fanout.split(",")]
val_fanouts = [int(fanout) for fanout in args.validation_fanout.split(",")]
sampler = dgl.dataloading.MultiLayerNeighborSampler(fanouts)
dataloader = dgl.distributed.DistNodeDataLoader(
g,
{"paper": train_nid},
sampler,
batch_size=args.batch_size,
shuffle=True,
drop_last=False,
)
valid_sampler = dgl.dataloading.MultiLayerNeighborSampler(val_fanouts)
valid_dataloader = dgl.distributed.DistNodeDataLoader(
g,
{"paper": val_nid},
valid_sampler,
batch_size=args.batch_size,
shuffle=False,
drop_last=False,
)
test_sampler = dgl.dataloading.MultiLayerNeighborSampler(val_fanouts)
test_dataloader = dgl.distributed.DistNodeDataLoader(
g,
{"paper": test_nid},
test_sampler,
batch_size=args.eval_batch_size,
shuffle=False,
drop_last=False,
)
embed_layer = DistEmbedLayer(
device,
g,
args.n_hidden,
sparse_emb=args.sparse_embedding,
dgl_sparse_emb=args.dgl_sparse,
feat_name="feat",
)
model = EntityClassify(
device,
args.n_hidden,
num_classes,
g.etypes,
num_bases=args.n_bases,
num_hidden_layers=args.n_layers - 2,
dropout=args.dropout,
use_self_loop=args.use_self_loop,
layer_norm=args.layer_norm,
)
model = model.to(device)
if not args.standalone:
if args.num_gpus == -1:
model = DistributedDataParallel(model)
# If there are dense parameters in the embedding layer
# or we use Pytorch saprse embeddings.
if len(embed_layer.node_projs) > 0 or not args.dgl_sparse:
embed_layer = DistributedDataParallel(embed_layer)
else:
dev_id = g.rank() % args.num_gpus
model = DistributedDataParallel(
model, device_ids=[dev_id], output_device=dev_id
)
# If there are dense parameters in the embedding layer
# or we use Pytorch saprse embeddings.
if len(embed_layer.node_projs) > 0 or not args.dgl_sparse:
embed_layer = embed_layer.to(device)
embed_layer = DistributedDataParallel(
embed_layer, device_ids=[dev_id], output_device=dev_id
)
if args.sparse_embedding:
if args.dgl_sparse and args.standalone:
emb_optimizer = dgl.distributed.optim.SparseAdam(
list(embed_layer.node_embeds.values()), lr=args.sparse_lr
)
print(
"optimize DGL sparse embedding:", embed_layer.node_embeds.keys()
)
elif args.dgl_sparse:
emb_optimizer = dgl.distributed.optim.SparseAdam(
list(embed_layer.module.node_embeds.values()), lr=args.sparse_lr
)
print(
"optimize DGL sparse embedding:",
embed_layer.module.node_embeds.keys(),
)
elif args.standalone:
emb_optimizer = th.optim.SparseAdam(
list(embed_layer.node_embeds.parameters()), lr=args.sparse_lr
)
print("optimize Pytorch sparse embedding:", embed_layer.node_embeds)
else:
emb_optimizer = th.optim.SparseAdam(
list(embed_layer.module.node_embeds.parameters()),
lr=args.sparse_lr,
)
print(
"optimize Pytorch sparse embedding:",
embed_layer.module.node_embeds,
)
dense_params = list(model.parameters())
if args.standalone:
dense_params += list(embed_layer.node_projs.parameters())
print("optimize dense projection:", embed_layer.node_projs)
else:
dense_params += list(embed_layer.module.node_projs.parameters())
print("optimize dense projection:", embed_layer.module.node_projs)
optimizer = th.optim.Adam(
dense_params, lr=args.lr, weight_decay=args.l2norm
)
else:
all_params = list(model.parameters()) + list(embed_layer.parameters())
optimizer = th.optim.Adam(
all_params, lr=args.lr, weight_decay=args.l2norm
)
# training loop
print("start training...")
for epoch in range(args.n_epochs):
tic = time.time()
sample_time = 0
copy_time = 0
forward_time = 0
backward_time = 0
update_time = 0
number_train = 0
number_input = 0
step_time = []
iter_t = []
sample_t = []
feat_copy_t = []
forward_t = []
backward_t = []
update_t = []
iter_tput = []
start = time.time()
# Loop over the dataloader to sample the computation dependency graph as a list of
# blocks.
step_time = []
for step, sample_data in enumerate(dataloader):
input_nodes, seeds, blocks = sample_data
seeds = seeds["paper"]
number_train += seeds.shape[0]
number_input += np.sum(
[blocks[0].num_src_nodes(ntype) for ntype in blocks[0].ntypes]
)
tic_step = time.time()
sample_time += tic_step - start
sample_t.append(tic_step - start)
feats = embed_layer(input_nodes)
label = labels[seeds].to(device)
copy_time = time.time()
feat_copy_t.append(copy_time - tic_step)
# forward
logits = model(blocks, feats)
assert len(logits) == 1
logits = logits["paper"]
loss = F.cross_entropy(logits, label)
forward_end = time.time()
# backward
optimizer.zero_grad()
if args.sparse_embedding:
emb_optimizer.zero_grad()
loss.backward()
compute_end = time.time()
forward_t.append(forward_end - copy_time)
backward_t.append(compute_end - forward_end)
# Update model parameters
optimizer.step()
if args.sparse_embedding:
emb_optimizer.step()
update_t.append(time.time() - compute_end)
step_t = time.time() - start
step_time.append(step_t)
train_acc = th.sum(logits.argmax(dim=1) == label).item() / len(
seeds
)
if step % args.log_every == 0:
print(
"[{}] Epoch {:05d} | Step {:05d} | Train acc {:.4f} | Loss {:.4f} | time {:.3f} s"
"| sample {:.3f} | copy {:.3f} | forward {:.3f} | backward {:.3f} | update {:.3f}".format(
g.rank(),
epoch,
step,
train_acc,
loss.item(),
np.sum(step_time[-args.log_every :]),
np.sum(sample_t[-args.log_every :]),
np.sum(feat_copy_t[-args.log_every :]),
np.sum(forward_t[-args.log_every :]),
np.sum(backward_t[-args.log_every :]),
np.sum(update_t[-args.log_every :]),
)
)
start = time.time()
gc.collect()
print(
"[{}]Epoch Time(s): {:.4f}, sample: {:.4f}, data copy: {:.4f}, forward: {:.4f}, backward: {:.4f}, update: {:.4f}, #train: {}, #input: {}".format(
g.rank(),
np.sum(step_time),
np.sum(sample_t),
np.sum(feat_copy_t),
np.sum(forward_t),
np.sum(backward_t),
np.sum(update_t),
number_train,
number_input,
)
)
epoch += 1
start = time.time()
g.barrier()
val_acc, test_acc = evaluate(
g,
model,
embed_layer,
labels,
valid_dataloader,
test_dataloader,
all_val_nid,
all_test_nid,
)
if val_acc >= 0:
print(
"Val Acc {:.4f}, Test Acc {:.4f}, time: {:.4f}".format(
val_acc, test_acc, time.time() - start
)
)
def main(args):
dgl.distributed.initialize(args.ip_config, use_graphbolt=args.use_graphbolt)
if not args.standalone:
backend = "gloo" if args.num_gpus == -1 else "nccl"
if args.sparse_embedding and args.dgl_sparse:
# `nccl` is not fully supported in DistDGL's sparse optimizer.
backend = "gloo"
th.distributed.init_process_group(backend=backend)
g = dgl.distributed.DistGraph(args.graph_name, part_config=args.conf_path)
print("rank:", g.rank())
pb = g.get_partition_book()
if "trainer_id" in g.nodes["paper"].data:
train_nid = dgl.distributed.node_split(
g.nodes["paper"].data["train_mask"],
pb,
ntype="paper",
force_even=True,
node_trainer_ids=g.nodes["paper"].data["trainer_id"],
)
val_nid = dgl.distributed.node_split(
g.nodes["paper"].data["val_mask"],
pb,
ntype="paper",
force_even=True,
node_trainer_ids=g.nodes["paper"].data["trainer_id"],
)
test_nid = dgl.distributed.node_split(
g.nodes["paper"].data["test_mask"],
pb,
ntype="paper",
force_even=True,
node_trainer_ids=g.nodes["paper"].data["trainer_id"],
)
else:
train_nid = dgl.distributed.node_split(
g.nodes["paper"].data["train_mask"],
pb,
ntype="paper",
force_even=True,
)
val_nid = dgl.distributed.node_split(
g.nodes["paper"].data["val_mask"],
pb,
ntype="paper",
force_even=True,
)
test_nid = dgl.distributed.node_split(
g.nodes["paper"].data["test_mask"],
pb,
ntype="paper",
force_even=True,
)
local_nid = pb.partid2nids(pb.partid, "paper").detach().numpy()
print(
"part {}, train: {} (local: {}), val: {} (local: {}), test: {} (local: {})".format(
g.rank(),
len(train_nid),
len(np.intersect1d(train_nid.numpy(), local_nid)),
len(val_nid),
len(np.intersect1d(val_nid.numpy(), local_nid)),
len(test_nid),
len(np.intersect1d(test_nid.numpy(), local_nid)),
)
)
if args.num_gpus == -1:
device = th.device("cpu")
else:
dev_id = g.rank() % args.num_gpus
device = th.device("cuda:" + str(dev_id))
labels = g.nodes["paper"].data["labels"][np.arange(g.num_nodes("paper"))]
all_val_nid = th.LongTensor(
np.nonzero(
g.nodes["paper"].data["val_mask"][np.arange(g.num_nodes("paper"))]
)
).squeeze()
all_test_nid = th.LongTensor(
np.nonzero(
g.nodes["paper"].data["test_mask"][np.arange(g.num_nodes("paper"))]
)
).squeeze()
n_classes = len(th.unique(labels[labels >= 0]))
print("#classes:", n_classes)
run(
args,
device,
(
g,
n_classes,
train_nid,
val_nid,
test_nid,
labels,
all_val_nid,
all_test_nid,
),
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="RGCN")
# distributed training related
parser.add_argument("--graph-name", type=str, help="graph name")
parser.add_argument("--id", type=int, help="the partition id")
parser.add_argument(
"--ip-config", type=str, help="The file for IP configuration"
)
parser.add_argument(
"--conf-path", type=str, help="The path to the partition config file"
)
# rgcn related
parser.add_argument(
"--num_gpus",
type=int,
default=-1,
help="the number of GPU device. Use -1 for CPU training",
)
parser.add_argument(
"--dropout", type=float, default=0, help="dropout probability"
)
parser.add_argument(
"--n-hidden", type=int, default=16, help="number of hidden units"
)
parser.add_argument("--lr", type=float, default=1e-2, help="learning rate")
parser.add_argument(
"--sparse-lr", type=float, default=1e-2, help="sparse lr rate"
)
parser.add_argument(
"--n-bases",
type=int,
default=-1,
help="number of filter weight matrices, default: -1 [use all]",
)
parser.add_argument(
"--n-layers", type=int, default=2, help="number of propagation rounds"
)
parser.add_argument(
"-e",
"--n-epochs",
type=int,
default=50,
help="number of training epochs",
)
parser.add_argument(
"-d", "--dataset", type=str, required=True, help="dataset to use"
)
parser.add_argument("--l2norm", type=float, default=0, help="l2 norm coef")
parser.add_argument(
"--relabel",
default=False,
action="store_true",
help="remove untouched nodes and relabel",
)
parser.add_argument(
"--fanout",
type=str,
default="4, 4",
help="Fan-out of neighbor sampling.",
)
parser.add_argument(
"--validation-fanout",
type=str,
default=None,
help="Fan-out of neighbor sampling during validation.",
)
parser.add_argument(
"--use-self-loop",
default=False,
action="store_true",
help="include self feature as a special relation",
)
parser.add_argument(
"--batch-size", type=int, default=100, help="Mini-batch size. "
)
parser.add_argument(
"--eval-batch-size", type=int, default=128, help="Mini-batch size. "
)
parser.add_argument("--log-every", type=int, default=20)
parser.add_argument(
"--low-mem",
default=False,
action="store_true",
help="Whether use low mem RelGraphCov",
)
parser.add_argument(
"--sparse-embedding",
action="store_true",
help="Use sparse embedding for node embeddings.",
)
parser.add_argument(
"--dgl-sparse",
action="store_true",
help="Whether to use DGL sparse embedding",
)
parser.add_argument(
"--layer-norm",
default=False,
action="store_true",
help="Use layer norm",
)
parser.add_argument(
"--local_rank", type=int, help="get rank of the process"
)
parser.add_argument(
"--standalone", action="store_true", help="run in the standalone mode"
)
parser.add_argument(
"--use_graphbolt",
action="store_true",
help="Use GraphBolt for distributed train.",
)
args = parser.parse_args()
# if validation_fanout is None, set it with args.fanout
if args.validation_fanout is None:
args.validation_fanout = args.fanout
print(args)
main(args)
examples/distributed/rgcn/partition_graph.py
+137
View File
@@ -0,0 +1,137 @@
import argparse
import time
import dgl
import numpy as np
import torch as th
from ogb.nodeproppred import DglNodePropPredDataset
def load_ogb(dataset):
if dataset == "ogbn-mag":
dataset = DglNodePropPredDataset(name=dataset)
split_idx = dataset.get_idx_split()
train_idx = split_idx["train"]["paper"]
val_idx = split_idx["valid"]["paper"]
test_idx = split_idx["test"]["paper"]
hg_orig, labels = dataset[0]
subgs = {}
for etype in hg_orig.canonical_etypes:
u, v = hg_orig.all_edges(etype=etype)
subgs[etype] = (u, v)
subgs[(etype[2], "rev-" + etype[1], etype[0])] = (v, u)
hg = dgl.heterograph(subgs)
hg.nodes["paper"].data["feat"] = hg_orig.nodes["paper"].data["feat"]
paper_labels = labels["paper"].squeeze()
num_rels = len(hg.canonical_etypes)
num_of_ntype = len(hg.ntypes)
num_classes = dataset.num_classes
category = "paper"
print("Number of relations: {}".format(num_rels))
print("Number of class: {}".format(num_classes))
print("Number of train: {}".format(len(train_idx)))
print("Number of valid: {}".format(len(val_idx)))
print("Number of test: {}".format(len(test_idx)))
# get target category id
category_id = len(hg.ntypes)
for i, ntype in enumerate(hg.ntypes):
if ntype == category:
category_id = i
train_mask = th.zeros((hg.num_nodes("paper"),), dtype=th.bool)
train_mask[train_idx] = True
val_mask = th.zeros((hg.num_nodes("paper"),), dtype=th.bool)
val_mask[val_idx] = True
test_mask = th.zeros((hg.num_nodes("paper"),), dtype=th.bool)
test_mask[test_idx] = True
hg.nodes["paper"].data["train_mask"] = train_mask
hg.nodes["paper"].data["val_mask"] = val_mask
hg.nodes["paper"].data["test_mask"] = test_mask
hg.nodes["paper"].data["labels"] = paper_labels
return hg
else:
raise ("Do not support other ogbn datasets.")
if __name__ == "__main__":
argparser = argparse.ArgumentParser("Partition builtin graphs")
argparser.add_argument(
"--dataset", type=str, default="ogbn-mag", help="datasets: ogbn-mag"
)
argparser.add_argument(
"--num_parts", type=int, default=4, help="number of partitions"
)
argparser.add_argument(
"--part_method", type=str, default="metis", help="the partition method"
)
argparser.add_argument(
"--balance_train",
action="store_true",
help="balance the training size in each partition.",
)
argparser.add_argument(
"--undirected",
action="store_true",
help="turn the graph into an undirected graph.",
)
argparser.add_argument(
"--balance_edges",
action="store_true",
help="balance the number of edges in each partition.",
)
argparser.add_argument(
"--num_trainers_per_machine",
type=int,
default=1,
help="the number of trainers per machine. The trainer ids are stored\
in the node feature 'trainer_id'",
)
argparser.add_argument(
"--output",
type=str,
default="data",
help="Output path of partitioned graph.",
)
argparser.add_argument(
"--use_graphbolt",
action="store_true",
help="Use GraphBolt for distributed train.",
)
args = argparser.parse_args()
start = time.time()
g = load_ogb(args.dataset)
print(
"load {} takes {:.3f} seconds".format(args.dataset, time.time() - start)
)
print("|V|={}, |E|={}".format(g.num_nodes(), g.num_edges()))
print(
"train: {}, valid: {}, test: {}".format(
th.sum(g.nodes["paper"].data["train_mask"]),
th.sum(g.nodes["paper"].data["val_mask"]),
th.sum(g.nodes["paper"].data["test_mask"]),
)
)
if args.balance_train:
balance_ntypes = {"paper": g.nodes["paper"].data["train_mask"]}
else:
balance_ntypes = None
dgl.distributed.partition_graph(
g,
args.dataset,
args.num_parts,
args.output,
part_method=args.part_method,
balance_ntypes=balance_ntypes,
balance_edges=args.balance_edges,
num_trainers_per_machine=args.num_trainers_per_machine,
use_graphbolt=args.use_graphbolt,
)