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

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/pytorch/graphsaint/README.md
+170
View File
@@ -0,0 +1,170 @@
# GraphSAINT
This DGL example implements the paper: GraphSAINT: Graph Sampling Based Inductive Learning Method.
Paper link: https://arxiv.org/abs/1907.04931
Author's code: https://github.com/GraphSAINT/GraphSAINT
Contributor: Jiahang Li ([@ljh1064126026](https://github.com/ljh1064126026)) Tang Liu ([@lt610](https://github.com/lt610))
For built-in GraphSAINT subgraph samplers with online sampling, use `dgl.dataloading.SAINTSampler`.
## Dependencies
- Python 3.7.10
- PyTorch 1.8.1
- NumPy 1.19.2
- Scikit-learn 0.23.2
- DGL 0.7.1
## Dataset
All datasets used are provided by Author's [code](https://github.com/GraphSAINT/GraphSAINT). They are available in [Google Drive](https://drive.google.com/drive/folders/1zycmmDES39zVlbVCYs88JTJ1Wm5FbfLz) (alternatively, [Baidu Wangpan (code: f1ao)](https://pan.baidu.com/s/1SOb0SiSAXavwAcNqkttwcg#list/path=%2F)). Dataset summary("m" stands for multi-label binary classification, and "s" for single-label.):
| Dataset | Nodes | Edges | Degree | Feature | Classes |
| :-: | :-: | :-: | :-: | :-: | :-: |
| PPI | 14,755 | 225,270 | 15 | 50 | 121(m) |
| Flickr | 89,250 | 899,756 | 10 | 500 | 7(s) |
| Reddit | 232,965 | 11,606,919 | 50 | 602 | 41(s) |
| Yelp | 716,847 | 6,977,410 | 10 | 300 | 100 (m) |
| Amazon | 1,598,960 | 132,169,734 | 83 | 200 | 107 (m) |
Note that the PPI dataset here is different from DGL's built-in variant.
## Config
- The config file is `config.py`, which contains best configs for experiments below.
- Please refer to `sampler.py` to see explanations of some key parameters.
### Parameters
| **aggr** | **arch** | **dataset** | **dropout** |
| ------------------------------------------------------------ | ------------------------------------------------------------ | ------------------------------------------------------------ | ------------------------------------------------------------ |
| define how to aggregate embeddings of each node and its neighbors' embeddings ,which can be 'concat', 'mean'. The neighbors' embeddings are generated based on GCN | e.g. '1-1-0', means there're three layers, the first and the second layer employ message passing on the graph, then aggregate the embeddings of each node and its neighbors. The last layer only updates each node's embedding. The message passing mechanism comes from GCN | the name of dataset, which can be 'ppi', 'flickr', 'reddit', 'yelp', 'amazon' | the dropout of model used in train_sampling.py |
| **edge_budget** | **gpu** | **length** | **log_dir** |
| the expected number of edges in each subgraph, which is specified in the paper | -1 means cpu, otherwise 'cuda:gpu', e.g. if gpu=0, use 'cuda:0' | the length of each random walk | the directory storing logs |
| **lr** | **n_epochs** | **n_hidden** | **no_batch_norm** |
| learning rate | training epochs | hidden dimension | True if do NOT employ batch normalization in each layer |
| **node_budget** | **num_subg** | **num_roots** | **sampler** |
| the expected number of nodes in each subgraph, which is specified in the paper | the expected number of pre_sampled subgraphs | the number of roots to generate random walks | specify which sampler to use, which can be 'node', 'edge', 'rw', corresponding to node, edge, random walk sampler |
| **use_val** | **val_every** | **num_workers_sampler** | **num_subg_sampler** |
| True if use best model to test, which is stored by earlystop mechanism | validate per 'val_every' epochs | number of workers (processes) specified for internal dataloader in SAINTSampler, which is to pre-sample subgraphs | the maximal number of pre-sampled subgraphs |
| **batch_size_sampler** | **num_workers** | | |
| batch size of internal dataloader in SAINTSampler | number of workers (processes) specified for external dataloader in train_sampling.py, which is to sample subgraphs in training phase | | |
## Minibatch training
Run with following:
```bash
python train_sampling.py --task $task $online
# online sampling: e.g. python train_sampling.py --task ppi_n --online
# offline sampling: e.g. python train_sampling.py --task flickr_e
```
- `$task` includes `ppi_n, ppi_e, ppi_rw, flickr_n, flickr_e, flickr_rw, reddit_n, reddit_e, reddit_rw, yelp_n, yelp_e, yelp_rw, amazon_n, amazon_e, amazon_rw`. For example, `ppi_n` represents running experiments on dataset `ppi` with `node sampler`
- If `$online` is `--online`, we sample subgraphs on-the-fly in the training phase, while discarding pre-sampled subgraphs. If `$online` is empty, we utilize pre-sampled subgraphs in the training phase.
## Experiments
* Paper: results from the paper
* Running: results from experiments with the authors' code
* DGL: results from experiments with the DGL example. The experiment config comes from `config.py`. You can modify parameters in the `config.py` to see different performance of different setup.
> Note that we implement offline sampling and online sampling in training phase. Offline sampling means all subgraphs utilized in training phase come from pre-sampled subgraphs. Online sampling means we discard all pre-sampled subgraphs and re-sample new subgraphs in training phase.
> Note that the sampling method in the pre-sampling phase must be offline sampling.
### F1-micro
#### Random node sampler
| Method | PPI | Flickr | Reddit | Yelp | Amazon |
| --- | --- | --- | --- | --- | --- |
| Paper | 0.960±0.001 | 0.507±0.001 | 0.962±0.001 | 0.641±0.000 | 0.782±0.004 |
| Running | 0.9628 | 0.5077 | 0.9622 | 0.6393 | 0.7695 |
| DGL_offline | 0.9715 | 0.5024 | 0.9645 | 0.6457 | 0.8051 |
| DGL_online | 0.9730 | 0.5071 | 0.9645 | 0.6444 | 0.8014 |
#### Random edge sampler
| Method | PPI | Flickr | Reddit | Yelp | Amazon |
| --- | --- | --- | --- | --- | --- |
| Paper | 0.981±0.007 | 0.510±0.002 | 0.966±0.001 | 0.653±0.003 | 0.807±0.001 |
| Running | 0.9810 | 0.5066 | 0.9656 | 0.6531 | 0.8071 |
| DGL_offline | 0.9817 | 0.5077 | 0.9655 | 0.6530 | 0.8034 |
| DGL_online | 0.9815 | 0.5041 | 0.9653 | 0.6516 | 0.7756 |
#### Random walk sampler
| Method | PPI | Flickr | Reddit | Yelp | Amazon |
| --- | --- | --- | --- | --- | --- |
| Paper | 0.981±0.004 | 0.511±0.001 | 0.966±0.001 | 0.653±0.003 | 0.815±0.001 |
| Running | 0.9812 | 0.5104 | 0.9648 | 0.6527 | 0.8131 |
| DGL_offline | 0.9833 | 0.5027 | 0.9582 | 0.6514 | 0.8178 |
| DGL_online | 0.9820 | 0.5110 | 0.9572 | 0.6508 | 0.8157 |
### Sampling time
- Here sampling time includes consumed time of pre-sampling subgraphs and calculating normalization coefficients in the beginning.
#### Random node sampler
| Method | PPI | Flickr | Reddit | Yelp | Amazon |
| --- | --- | --- | --- | --- | --- |
| Running | 1.46 | 3.49 | 19 | 59.01 | 978.62 |
| DGL | 2.51 | 1.12 | 27.32 | 60.15 | 929.24 |
#### Random edge sampler
| Method | PPI | Flickr | Reddit | Yelp | Amazon |
| --- | --- | --- | --- | --- | --- |
| Running | 1.4 | 3.18 | 13.88 | 39.02 | |
| DGL | 3.04 | 1.87 | 52.01 | 48.38 | |
#### Random walk sampler
| Method | PPI | Flickr | Reddit | Yelp | Amazon |
| --- | --- | --- | --- | --- | --- |
| Running | 1.7 | 3.82 | 16.97 | 43.25 | 355.68 |
| DGL | 3.05 | 2.13 | 11.01 | 22.23 | 151.84 |
## Test std of sampling and normalization time
- We've run experiments 10 times repeatedly to test average and standard deviation of sampling and normalization time. Here we just test time without training model to the end. Moreover, for efficient testing, the hardware and config employed here are not the same as the experiments above, so the sampling time might be a bit different from that above. But we keep the environment consistent in all experiments below.
> The config here which is different with that in the section above is only `num_workers_sampler`, `batch_size_sampler` and `num_workers`, which are only correlated to the sampling speed. Other parameters are kept consistent across two sections thus the model's performance is not affected.
> The value is (average, std).
### Random node sampler
| Method | PPI | Flickr | Reddit | Yelp | Amazon |
| ------------------------- | --------------- | ------------ | ------------- | ------------- | --------------- |
| DGL_Sampling(std) | 2.618, 0.004 | 3.017, 0.507 | 35.356, 2.363 | 69.913, 6.3 | 888.025, 16.004 |
| DGL_Normalization(std) | Small to ignore | 0.008, 0.004 | 0.26, 0.047 | 0.189, 0.0288 | 2.443, 0.124 |
| | | | | | |
| author_Sampling(std) | 0.788, 0.661 | 0.728, 0.367 | 8.931, 3.155 | 27.818, 1.384 | 295.597, 4.928 |
| author_Normalization(std) | 0.665, 0.565 | 4.981, 2.952 | 17.231, 7.116 | 47.449, 2.794 | 279.241, 17.615 |
### Random edge sampler
| Method | PPI | Flickr | Reddit | Yelp | Amazon |
| ------------------------- | --------------- | ------------ | ------------- | ------------- | ------ |
| DGL_Sampling(std) | 3.554, 0.292 | 4.722, 0.245 | 47.09, 2.76 | 75.219, 6.442 | |
| DGL_Normalization(std) | Small to ignore | 0.005, 0.007 | 0.235, 0.026 | 0.193, 0.021 | |
| | | | | | |
| author_Sampling(std) | 0.802, 0.667 | 0.761, 0.387 | 6.058, 2.166 | 13.914, 1.864 | |
| author_Normalization(std) | 0.667, 0.570 | 5.180, 3.006 | 15.803, 5.867 | 44.278, 5.853 | |
### Random walk sampler
| Method | PPI | Flickr | Reddit | Yelp | Amazon |
| ------------------------- | --------------- | ------------ | ------------- | ------------- | --------------- |
| DGL_Sampling(std) | 3.304, 0.08 | 5.487, 1.294 | 37.041, 2.083 | 39.951, 3.094 | 179.613, 18.881 |
| DGL_Normalization(std) | Small to ignore | 0.001, 0.003 | 0.235, 0.026 | 0.185, 0.018 | 3.769, 0.326 |
| | | | | | |
| author_Sampling(std) | 0.924, 0.773 | 1.405, 0.718 | 8.608, 3.093 | 19.113, 1.700 | 217.184, 1.546 |
| author_Normalization(std) | 0.701, 0.596 | 5.025, 2.954 | 18.198, 7.223 | 45.874, 8.020 | 128.272, 3.170 |
examples/pytorch/graphsaint/config.py
+364
View File
@@ -0,0 +1,364 @@
CONFIG = {
"ppi_n": {
"aggr": "concat",
"arch": "1-0-1-0",
"dataset": "ppi",
"dropout": 0,
"edge_budget": 4000,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 50,
"n_hidden": 512,
"no_batch_norm": False,
"node_budget": 6000,
"num_subg": 50,
"num_roots": 3000,
"sampler": "node",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 0,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"ppi_e": {
"aggr": "concat",
"arch": "1-0-1-0",
"dataset": "ppi",
"dropout": 0.1,
"edge_budget": 4000,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 50,
"n_hidden": 512,
"no_batch_norm": False,
"node_budget": 6000,
"num_subg": 50,
"num_roots": 3000,
"sampler": "edge",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 0,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"ppi_rw": {
"aggr": "concat",
"arch": "1-0-1-0",
"dataset": "ppi",
"dropout": 0.1,
"edge_budget": 4000,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 50,
"n_hidden": 512,
"no_batch_norm": False,
"node_budget": 6000,
"num_subg": 50,
"num_roots": 3000,
"sampler": "rw",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 0,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"flickr_n": {
"aggr": "concat",
"arch": "1-1-0",
"dataset": "flickr",
"dropout": 0.2,
"edge_budget": 6000,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 50,
"n_hidden": 256,
"no_batch_norm": False,
"node_budget": 8000,
"num_subg": 25,
"num_roots": 6000,
"sampler": "node",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 0,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": False,
},
"flickr_e": {
"aggr": "concat",
"arch": "1-1-0",
"dataset": "flickr",
"dropout": 0.2,
"edge_budget": 6000,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 50,
"n_hidden": 256,
"no_batch_norm": False,
"node_budget": 8000,
"num_subg": 25,
"num_roots": 6000,
"sampler": "edge",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 0,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": False,
},
"flickr_rw": {
"aggr": "concat",
"arch": "1-1-0",
"dataset": "flickr",
"dropout": 0.2,
"edge_budget": 6000,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 50,
"n_hidden": 256,
"no_batch_norm": False,
"node_budget": 8000,
"num_subg": 25,
"num_roots": 6000,
"sampler": "rw",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 0,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": False,
},
"reddit_n": {
"aggr": "concat",
"arch": "1-0-1-0",
"dataset": "reddit",
"dropout": 0.1,
"edge_budget": 4000,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 20,
"n_hidden": 128,
"no_batch_norm": False,
"node_budget": 8000,
"num_subg": 50,
"num_roots": 3000,
"sampler": "node",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 8,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"reddit_e": {
"aggr": "concat",
"arch": "1-0-1-0",
"dataset": "reddit",
"dropout": 0.1,
"edge_budget": 6000,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 20,
"n_hidden": 128,
"no_batch_norm": False,
"node_budget": 8000,
"num_subg": 50,
"num_roots": 3000,
"sampler": "edge",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 8,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"reddit_rw": {
"aggr": "concat",
"arch": "1-0-1-0",
"dataset": "reddit",
"dropout": 0.1,
"edge_budget": 6000,
"length": 4,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 10,
"n_hidden": 128,
"no_batch_norm": False,
"node_budget": 8000,
"num_subg": 50,
"num_roots": 200,
"sampler": "rw",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 8,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"yelp_n": {
"aggr": "concat",
"arch": "1-1-0",
"dataset": "yelp",
"dropout": 0.1,
"edge_budget": 6000,
"length": 4,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 10,
"n_hidden": 512,
"no_batch_norm": False,
"node_budget": 5000,
"num_subg": 50,
"num_roots": 200,
"sampler": "node",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 8,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"yelp_e": {
"aggr": "concat",
"arch": "1-1-0",
"dataset": "yelp",
"dropout": 0.1,
"edge_budget": 2500,
"length": 4,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 10,
"n_hidden": 512,
"no_batch_norm": False,
"node_budget": 5000,
"num_subg": 50,
"num_roots": 200,
"sampler": "edge",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 8,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"yelp_rw": {
"aggr": "concat",
"arch": "1-1-0",
"dataset": "yelp",
"dropout": 0.1,
"edge_budget": 2500,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 10,
"n_hidden": 512,
"no_batch_norm": False,
"node_budget": 5000,
"num_subg": 50,
"num_roots": 1250,
"sampler": "rw",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 8,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"amazon_n": {
"aggr": "concat",
"arch": "1-1-0",
"dataset": "amazon",
"dropout": 0.1,
"edge_budget": 2500,
"length": 4,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 5,
"n_hidden": 512,
"no_batch_norm": False,
"node_budget": 4500,
"num_subg": 50,
"num_roots": 200,
"sampler": "node",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 4,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
"amazon_e": {
"aggr": "concat",
"arch": "1-1-0",
"dataset": "amazon",
"dropout": 0.1,
"edge_budget": 2000,
"gpu": 0,
"length": 4,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 10,
"n_hidden": 512,
"no_batch_norm": False,
"node_budget": 5000,
"num_subg": 50,
"num_roots": 200,
"sampler": "edge",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 20,
"num_subg_sampler": 5000,
"batch_size_sampler": 50,
"num_workers": 26,
"full": True,
},
"amazon_rw": {
"aggr": "concat",
"arch": "1-1-0",
"dataset": "amazon",
"dropout": 0.1,
"edge_budget": 2500,
"gpu": 0,
"length": 2,
"log_dir": "none",
"lr": 0.01,
"n_epochs": 5,
"n_hidden": 512,
"no_batch_norm": False,
"node_budget": 5000,
"num_subg": 50,
"num_roots": 1500,
"sampler": "rw",
"use_val": True,
"val_every": 1,
"num_workers_sampler": 4,
"num_subg_sampler": 10000,
"batch_size_sampler": 200,
"num_workers": 8,
"full": True,
},
}
examples/pytorch/graphsaint/modules.py
+165
View File
@@ -0,0 +1,165 @@
import dgl.function as fn
import torch as th
import torch.nn as nn
import torch.nn.functional as F
class GCNLayer(nn.Module):
def __init__(
self,
in_dim,
out_dim,
order=1,
act=None,
dropout=0,
batch_norm=False,
aggr="concat",
):
super(GCNLayer, self).__init__()
self.lins = nn.ModuleList()
self.bias = nn.ParameterList()
for _ in range(order + 1):
self.lins.append(nn.Linear(in_dim, out_dim, bias=False))
self.bias.append(nn.Parameter(th.zeros(out_dim)))
self.order = order
self.act = act
self.dropout = nn.Dropout(dropout)
self.batch_norm = batch_norm
if batch_norm:
self.offset, self.scale = nn.ParameterList(), nn.ParameterList()
for _ in range(order + 1):
self.offset.append(nn.Parameter(th.zeros(out_dim)))
self.scale.append(nn.Parameter(th.ones(out_dim)))
self.aggr = aggr
self.reset_parameters()
def reset_parameters(self):
for lin in self.lins:
nn.init.xavier_normal_(lin.weight)
def feat_trans(
self, features, idx
): # linear transformation + activation + batch normalization
h = self.lins[idx](features) + self.bias[idx]
if self.act is not None:
h = self.act(h)
if self.batch_norm:
mean = h.mean(dim=1).view(h.shape[0], 1)
var = h.var(dim=1, unbiased=False).view(h.shape[0], 1) + 1e-9
h = (h - mean) * self.scale[idx] * th.rsqrt(var) + self.offset[idx]
return h
def forward(self, graph, features):
g = graph.local_var()
h_in = self.dropout(features)
h_hop = [h_in]
D_norm = (
g.ndata["train_D_norm"]
if "train_D_norm" in g.ndata
else g.ndata["full_D_norm"]
)
for _ in range(self.order): # forward propagation
g.ndata["h"] = h_hop[-1]
if "w" not in g.edata:
g.edata["w"] = th.ones((g.num_edges(),)).to(features.device)
g.update_all(fn.u_mul_e("h", "w", "m"), fn.sum("m", "h"))
h = g.ndata.pop("h")
h = h * D_norm
h_hop.append(h)
h_part = [self.feat_trans(ft, idx) for idx, ft in enumerate(h_hop)]
if self.aggr == "mean":
h_out = h_part[0]
for i in range(len(h_part) - 1):
h_out = h_out + h_part[i + 1]
elif self.aggr == "concat":
h_out = th.cat(h_part, 1)
else:
raise NotImplementedError
return h_out
class GCNNet(nn.Module):
def __init__(
self,
in_dim,
hid_dim,
out_dim,
arch="1-1-0",
act=F.relu,
dropout=0,
batch_norm=False,
aggr="concat",
):
super(GCNNet, self).__init__()
self.gcn = nn.ModuleList()
orders = list(map(int, arch.split("-")))
self.gcn.append(
GCNLayer(
in_dim=in_dim,
out_dim=hid_dim,
order=orders[0],
act=act,
dropout=dropout,
batch_norm=batch_norm,
aggr=aggr,
)
)
pre_out = ((aggr == "concat") * orders[0] + 1) * hid_dim
for i in range(1, len(orders) - 1):
self.gcn.append(
GCNLayer(
in_dim=pre_out,
out_dim=hid_dim,
order=orders[i],
act=act,
dropout=dropout,
batch_norm=batch_norm,
aggr=aggr,
)
)
pre_out = ((aggr == "concat") * orders[i] + 1) * hid_dim
self.gcn.append(
GCNLayer(
in_dim=pre_out,
out_dim=hid_dim,
order=orders[-1],
act=act,
dropout=dropout,
batch_norm=batch_norm,
aggr=aggr,
)
)
pre_out = ((aggr == "concat") * orders[-1] + 1) * hid_dim
self.out_layer = GCNLayer(
in_dim=pre_out,
out_dim=out_dim,
order=0,
act=None,
dropout=dropout,
batch_norm=False,
aggr=aggr,
)
def forward(self, graph):
h = graph.ndata["feat"]
for layer in self.gcn:
h = layer(graph, h)
h = F.normalize(h, p=2, dim=1)
h = self.out_layer(graph, h)
return h
examples/pytorch/graphsaint/sampler.py
+395
View File
@@ -0,0 +1,395 @@
import math
import os
import random
import time
import dgl
import dgl.function as fn
import numpy as np
import scipy
import torch as th
from dgl.sampling import pack_traces, random_walk
from torch.utils.data import DataLoader
# The base class of sampler
class SAINTSampler:
"""
Description
-----------
SAINTSampler implements the sampler described in GraphSAINT. This sampler implements offline sampling in
pre-sampling phase as well as fully offline sampling, fully online sampling in training phase.
Users can conveniently set param 'online' of the sampler to choose different modes.
Parameters
----------
node_budget : int
the expected number of nodes in each subgraph, which is specifically explained in the paper. Actually this
param specifies the times of sampling nodes from the original graph with replacement. The meaning of edge_budget
is similar to the node_budget.
dn : str
name of dataset.
g : DGLGraph
the full graph.
train_nid : list
ids of training nodes.
num_workers_sampler : int
number of processes to sample subgraphs in pre-sampling procedure using torch.dataloader.
num_subg_sampler : int, optional
the max number of subgraphs sampled in pre-sampling phase for computing normalization coefficients in the beginning.
Actually this param is used as ``__len__`` of sampler in pre-sampling phase.
Please make sure that num_subg_sampler is greater than batch_size_sampler so that we can sample enough subgraphs.
Defaults: 10000
batch_size_sampler : int, optional
the number of subgraphs sampled by each process concurrently in pre-sampling phase.
Defaults: 200
online : bool, optional
If `True`, we employ online sampling in training phase. Otherwise employing offline sampling.
Defaults: True
num_subg : int, optional
the expected number of sampled subgraphs in pre-sampling phase.
It is actually the 'N' in the original paper. Note that this param is different from the num_subg_sampler.
This param is just used to control the number of pre-sampled subgraphs.
Defaults: 50
full : bool, optional
True if the number of subgraphs used in the training phase equals to that of pre-sampled subgraphs, or
``math.ceil(self.train_g.num_nodes() / self.node_budget)``. This formula takes the result of A divided by B as
the number of subgraphs used in the training phase, where A is the number of training nodes in the original
graph, B is the expected number of nodes in each pre-sampled subgraph. Please refer to the paper to check the
details.
Defaults: True
Notes
-----
For parallelism of pre-sampling, we utilize `torch.DataLoader` to concurrently speed up sampling.
The `num_subg_sampler` is the return value of `__len__` in pre-sampling phase. Moreover, the param `batch_size_sampler`
determines the batch_size of `torch.DataLoader` in internal pre-sampling part. But note that if we wanna pass the
SAINTSampler to `torch.DataLoader` for concurrently sampling subgraphs in training phase, we need to specify
`batch_size` of `DataLoader`, that is, `batch_size_sampler` is not related to how sampler works in training procedure.
"""
def __init__(
self,
node_budget,
dn,
g,
train_nid,
num_workers_sampler,
num_subg_sampler=10000,
batch_size_sampler=200,
online=True,
num_subg=50,
full=True,
):
self.g = g.cpu()
self.node_budget = node_budget
self.train_g: dgl.graph = g.subgraph(train_nid)
self.dn, self.num_subg = dn, num_subg
self.node_counter = th.zeros((self.train_g.num_nodes(),))
self.edge_counter = th.zeros((self.train_g.num_edges(),))
self.prob = None
self.num_subg_sampler = num_subg_sampler
self.batch_size_sampler = batch_size_sampler
self.num_workers_sampler = num_workers_sampler
self.train = False
self.online = online
self.full = full
assert (
self.num_subg_sampler >= self.batch_size_sampler
), "num_subg_sampler should be greater than batch_size_sampler"
graph_fn, norm_fn = self.__generate_fn__()
if os.path.exists(graph_fn):
self.subgraphs = np.load(graph_fn, allow_pickle=True)
aggr_norm, loss_norm = np.load(norm_fn, allow_pickle=True)
else:
os.makedirs("./subgraphs/", exist_ok=True)
self.subgraphs = []
self.N, sampled_nodes = 0, 0
# N: the number of pre-sampled subgraphs
# Employ parallelism to speed up the sampling procedure
loader = DataLoader(
self,
batch_size=self.batch_size_sampler,
shuffle=True,
num_workers=self.num_workers_sampler,
collate_fn=self.__collate_fn__,
drop_last=False,
)
t = time.perf_counter()
for num_nodes, subgraphs_nids, subgraphs_eids in loader:
self.subgraphs.extend(subgraphs_nids)
sampled_nodes += num_nodes
_subgraphs, _node_counts = np.unique(
np.concatenate(subgraphs_nids), return_counts=True
)
sampled_nodes_idx = th.from_numpy(_subgraphs)
_node_counts = th.from_numpy(_node_counts)
self.node_counter[sampled_nodes_idx] += _node_counts
_subgraphs_eids, _edge_counts = np.unique(
np.concatenate(subgraphs_eids), return_counts=True
)
sampled_edges_idx = th.from_numpy(_subgraphs_eids)
_edge_counts = th.from_numpy(_edge_counts)
self.edge_counter[sampled_edges_idx] += _edge_counts
self.N += len(subgraphs_nids) # number of subgraphs
if sampled_nodes > self.train_g.num_nodes() * num_subg:
break
print(f"Sampling time: [{time.perf_counter() - t:.2f}s]")
np.save(graph_fn, self.subgraphs)
t = time.perf_counter()
aggr_norm, loss_norm = self.__compute_norm__()
print(f"Normalization time: [{time.perf_counter() - t:.2f}s]")
np.save(norm_fn, (aggr_norm, loss_norm))
self.train_g.ndata["l_n"] = th.Tensor(loss_norm)
self.train_g.edata["w"] = th.Tensor(aggr_norm)
self.__compute_degree_norm() # basically normalizing adjacent matrix
random.shuffle(self.subgraphs)
self.__clear__()
print("The number of subgraphs is: ", len(self.subgraphs))
self.train = True
def __len__(self):
if self.train is False:
return self.num_subg_sampler
else:
if self.full:
return len(self.subgraphs)
else:
return math.ceil(self.train_g.num_nodes() / self.node_budget)
def __getitem__(self, idx):
# Only when sampling subgraphs in training procedure and need to utilize sampled subgraphs and we still
# have sampled subgraphs we can fetch a subgraph from sampled subgraphs
if self.train:
if self.online:
subgraph = self.__sample__()
return dgl.node_subgraph(self.train_g, subgraph)
else:
return dgl.node_subgraph(self.train_g, self.subgraphs[idx])
else:
subgraph_nids = self.__sample__()
num_nodes = len(subgraph_nids)
subgraph_eids = dgl.node_subgraph(
self.train_g, subgraph_nids
).edata[dgl.EID]
return num_nodes, subgraph_nids, subgraph_eids
def __collate_fn__(self, batch):
if (
self.train
): # sample only one graph each epoch, batch_size in training phase in 1
return batch[0]
else:
sum_num_nodes = 0
subgraphs_nids_list = []
subgraphs_eids_list = []
for num_nodes, subgraph_nids, subgraph_eids in batch:
sum_num_nodes += num_nodes
subgraphs_nids_list.append(subgraph_nids)
subgraphs_eids_list.append(subgraph_eids)
return sum_num_nodes, subgraphs_nids_list, subgraphs_eids_list
def __clear__(self):
self.prob = None
self.node_counter = None
self.edge_counter = None
self.g = None
def __generate_fn__(self):
raise NotImplementedError
def __compute_norm__(self):
self.node_counter[self.node_counter == 0] = 1
self.edge_counter[self.edge_counter == 0] = 1
loss_norm = self.N / self.node_counter / self.train_g.num_nodes()
self.train_g.ndata["n_c"] = self.node_counter
self.train_g.edata["e_c"] = self.edge_counter
self.train_g.apply_edges(fn.v_div_e("n_c", "e_c", "a_n"))
aggr_norm = self.train_g.edata.pop("a_n")
self.train_g.ndata.pop("n_c")
self.train_g.edata.pop("e_c")
return aggr_norm.numpy(), loss_norm.numpy()
def __compute_degree_norm(self):
self.train_g.ndata[
"train_D_norm"
] = 1.0 / self.train_g.in_degrees().float().clamp(min=1).unsqueeze(1)
self.g.ndata["full_D_norm"] = 1.0 / self.g.in_degrees().float().clamp(
min=1
).unsqueeze(1)
def __sample__(self):
raise NotImplementedError
class SAINTNodeSampler(SAINTSampler):
"""
Description
-----------
GraphSAINT with node sampler.
Parameters
----------
node_budget : int
the expected number of nodes in each subgraph, which is specifically explained in the paper.
"""
def __init__(self, node_budget, **kwargs):
self.node_budget = node_budget
super(SAINTNodeSampler, self).__init__(
node_budget=node_budget, **kwargs
)
def __generate_fn__(self):
graph_fn = os.path.join(
"./subgraphs/{}_Node_{}_{}.npy".format(
self.dn, self.node_budget, self.num_subg
)
)
norm_fn = os.path.join(
"./subgraphs/{}_Node_{}_{}_norm.npy".format(
self.dn, self.node_budget, self.num_subg
)
)
return graph_fn, norm_fn
def __sample__(self):
if self.prob is None:
self.prob = self.train_g.in_degrees().float().clamp(min=1)
sampled_nodes = th.multinomial(
self.prob, num_samples=self.node_budget, replacement=True
).unique()
return sampled_nodes.numpy()
class SAINTEdgeSampler(SAINTSampler):
"""
Description
-----------
GraphSAINT with edge sampler.
Parameters
----------
edge_budget : int
the expected number of edges in each subgraph, which is specifically explained in the paper.
"""
def __init__(self, edge_budget, **kwargs):
self.edge_budget = edge_budget
self.rng = np.random.default_rng()
super(SAINTEdgeSampler, self).__init__(
node_budget=edge_budget * 2, **kwargs
)
def __generate_fn__(self):
graph_fn = os.path.join(
"./subgraphs/{}_Edge_{}_{}.npy".format(
self.dn, self.edge_budget, self.num_subg
)
)
norm_fn = os.path.join(
"./subgraphs/{}_Edge_{}_{}_norm.npy".format(
self.dn, self.edge_budget, self.num_subg
)
)
return graph_fn, norm_fn
# TODO: only sample half edges, then add another half edges
# TODO: use numpy to implement cython sampling method
def __sample__(self):
if self.prob is None:
src, dst = self.train_g.edges()
src_degrees, dst_degrees = self.train_g.in_degrees(
src
).float().clamp(min=1), self.train_g.in_degrees(dst).float().clamp(
min=1
)
prob_mat = 1.0 / src_degrees + 1.0 / dst_degrees
prob_mat = scipy.sparse.csr_matrix(
(prob_mat.numpy(), (src.numpy(), dst.numpy()))
)
# The edge probability here only contains that of edges in upper triangle adjacency matrix
# Because we assume the graph is undirected, that is, the adjacency matrix is symmetric. We only need
# to consider half of edges in the graph.
self.prob = th.tensor(scipy.sparse.triu(prob_mat).data)
self.prob /= self.prob.sum()
self.adj_nodes = np.stack(prob_mat.nonzero(), axis=1)
sampled_edges = np.unique(
dgl.random.choice(
len(self.prob),
size=self.edge_budget,
prob=self.prob,
replace=False,
)
)
sampled_nodes = np.unique(
self.adj_nodes[sampled_edges].flatten()
).astype("long")
return sampled_nodes
class SAINTRandomWalkSampler(SAINTSampler):
"""
Description
-----------
GraphSAINT with random walk sampler
Parameters
----------
num_roots : int
the number of roots to generate random walks.
length : int
the length of each random walk.
"""
def __init__(self, num_roots, length, **kwargs):
self.num_roots, self.length = num_roots, length
super(SAINTRandomWalkSampler, self).__init__(
node_budget=num_roots * length, **kwargs
)
def __generate_fn__(self):
graph_fn = os.path.join(
"./subgraphs/{}_RW_{}_{}_{}.npy".format(
self.dn, self.num_roots, self.length, self.num_subg
)
)
norm_fn = os.path.join(
"./subgraphs/{}_RW_{}_{}_{}_norm.npy".format(
self.dn, self.num_roots, self.length, self.num_subg
)
)
return graph_fn, norm_fn
def __sample__(self):
sampled_roots = th.randint(
0, self.train_g.num_nodes(), (self.num_roots,)
)
traces, types = random_walk(
self.train_g, nodes=sampled_roots, length=self.length
)
sampled_nodes, _, _, _ = pack_traces(traces, types)
sampled_nodes = sampled_nodes.unique()
return sampled_nodes.numpy()
examples/pytorch/graphsaint/train_sampling.py
+252
View File
@@ -0,0 +1,252 @@
import argparse
import os
import time
import warnings
import torch
import torch.nn.functional as F
from config import CONFIG
from modules import GCNNet
from sampler import SAINTEdgeSampler, SAINTNodeSampler, SAINTRandomWalkSampler
from torch.utils.data import DataLoader
from utils import calc_f1, evaluate, load_data, Logger, save_log_dir
def main(args, task):
warnings.filterwarnings("ignore")
multilabel_data = {"ppi", "yelp", "amazon"}
multilabel = args.dataset in multilabel_data
# This flag is excluded for too large dataset, like amazon, the graph of which is too large to be directly
# shifted to one gpu. So we need to
# 1. put the whole graph on cpu, and put the subgraphs on gpu in training phase
# 2. put the model on gpu in training phase, and put the model on cpu in validation/testing phase
# We need to judge cpu_flag and cuda (below) simultaneously when shift model between cpu and gpu
if args.dataset in ["amazon"]:
cpu_flag = True
else:
cpu_flag = False
# load and preprocess dataset
data = load_data(args, multilabel)
g = data.g
train_mask = g.ndata["train_mask"]
val_mask = g.ndata["val_mask"]
test_mask = g.ndata["test_mask"]
labels = g.ndata["label"]
train_nid = data.train_nid
in_feats = g.ndata["feat"].shape[1]
n_classes = data.num_classes
n_nodes = g.num_nodes()
n_edges = g.num_edges()
n_train_samples = train_mask.int().sum().item()
n_val_samples = val_mask.int().sum().item()
n_test_samples = test_mask.int().sum().item()
print(
"""----Data statistics------'
#Nodes %d
#Edges %d
#Classes/Labels (multi binary labels) %d
#Train samples %d
#Val samples %d
#Test samples %d"""
% (
n_nodes,
n_edges,
n_classes,
n_train_samples,
n_val_samples,
n_test_samples,
)
)
# load sampler
kwargs = {
"dn": args.dataset,
"g": g,
"train_nid": train_nid,
"num_workers_sampler": args.num_workers_sampler,
"num_subg_sampler": args.num_subg_sampler,
"batch_size_sampler": args.batch_size_sampler,
"online": args.online,
"num_subg": args.num_subg,
"full": args.full,
}
if args.sampler == "node":
saint_sampler = SAINTNodeSampler(args.node_budget, **kwargs)
elif args.sampler == "edge":
saint_sampler = SAINTEdgeSampler(args.edge_budget, **kwargs)
elif args.sampler == "rw":
saint_sampler = SAINTRandomWalkSampler(
args.num_roots, args.length, **kwargs
)
else:
raise NotImplementedError
loader = DataLoader(
saint_sampler,
collate_fn=saint_sampler.__collate_fn__,
batch_size=1,
shuffle=True,
num_workers=args.num_workers,
drop_last=False,
)
# set device for dataset tensors
if args.gpu < 0:
cuda = False
else:
cuda = True
torch.cuda.set_device(args.gpu)
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
if not cpu_flag:
g = g.to("cuda:{}".format(args.gpu))
print("labels shape:", g.ndata["label"].shape)
print("features shape:", g.ndata["feat"].shape)
model = GCNNet(
in_dim=in_feats,
hid_dim=args.n_hidden,
out_dim=n_classes,
arch=args.arch,
dropout=args.dropout,
batch_norm=not args.no_batch_norm,
aggr=args.aggr,
)
if cuda:
model.cuda()
# logger and so on
log_dir = save_log_dir(args)
logger = Logger(os.path.join(log_dir, "loggings"))
logger.write(args)
# use optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
# set train_nids to cuda tensor
if cuda:
train_nid = torch.from_numpy(train_nid).cuda()
print(
"GPU memory allocated before training(MB)",
torch.cuda.memory_allocated(device=train_nid.device) / 1024 / 1024,
)
start_time = time.time()
best_f1 = -1
for epoch in range(args.n_epochs):
for j, subg in enumerate(loader):
if cuda:
subg = subg.to(torch.cuda.current_device())
model.train()
# forward
pred = model(subg)
batch_labels = subg.ndata["label"]
if multilabel:
loss = F.binary_cross_entropy_with_logits(
pred,
batch_labels,
reduction="sum",
weight=subg.ndata["l_n"].unsqueeze(1),
)
else:
loss = F.cross_entropy(pred, batch_labels, reduction="none")
loss = (subg.ndata["l_n"] * loss).sum()
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm(model.parameters(), 5)
optimizer.step()
if j == len(loader) - 1:
model.eval()
with torch.no_grad():
train_f1_mic, train_f1_mac = calc_f1(
batch_labels.cpu().numpy(),
pred.cpu().numpy(),
multilabel,
)
print(
f"epoch:{epoch + 1}/{args.n_epochs}, Iteration {j + 1}/"
f"{len(loader)}:training loss",
loss.item(),
)
print(
"Train F1-mic {:.4f}, Train F1-mac {:.4f}".format(
train_f1_mic, train_f1_mac
)
)
# evaluate
model.eval()
if epoch % args.val_every == 0:
if (
cpu_flag and cuda
): # Only when we have shifted model to gpu and we need to shift it back on cpu
model = model.to("cpu")
val_f1_mic, val_f1_mac = evaluate(
model, g, labels, val_mask, multilabel
)
print(
"Val F1-mic {:.4f}, Val F1-mac {:.4f}".format(
val_f1_mic, val_f1_mac
)
)
if val_f1_mic > best_f1:
best_f1 = val_f1_mic
print("new best val f1:", best_f1)
torch.save(
model.state_dict(),
os.path.join(log_dir, "best_model_{}.pkl".format(task)),
)
if cpu_flag and cuda:
model.cuda()
end_time = time.time()
print(f"training using time {end_time - start_time}")
# test
if args.use_val:
model.load_state_dict(
torch.load(
os.path.join(log_dir, "best_model_{}.pkl".format(task)),
weights_only=False,
)
)
if cpu_flag and cuda:
model = model.to("cpu")
test_f1_mic, test_f1_mac = evaluate(model, g, labels, test_mask, multilabel)
print(
"Test F1-mic {:.4f}, Test F1-mac {:.4f}".format(
test_f1_mic, test_f1_mac
)
)
if __name__ == "__main__":
warnings.filterwarnings("ignore")
parser = argparse.ArgumentParser(description="GraphSAINT")
parser.add_argument(
"--task", type=str, default="ppi_n", help="type of tasks"
)
parser.add_argument(
"--online",
dest="online",
action="store_true",
help="sampling method in training phase",
)
parser.add_argument("--gpu", type=int, default=0, help="the gpu index")
task = parser.parse_args().task
args = argparse.Namespace(**CONFIG[task])
args.online = parser.parse_args().online
args.gpu = parser.parse_args().gpu
print(args)
main(args, task=task)
examples/pytorch/graphsaint/utils.py
+120
View File
@@ -0,0 +1,120 @@
import json
import os
from functools import namedtuple
import dgl
import numpy as np
import scipy.sparse
import torch
from sklearn.metrics import f1_score
from sklearn.preprocessing import StandardScaler
class Logger(object):
"""A custom logger to log stdout to a logging file."""
def __init__(self, path):
"""Initialize the logger.
Parameters
---------
path : str
The file path to be stored in.
"""
self.path = path
def write(self, s):
with open(self.path, "a") as f:
f.write(str(s))
print(s)
return
def save_log_dir(args):
log_dir = "./log/{}/{}".format(args.dataset, args.log_dir)
os.makedirs(log_dir, exist_ok=True)
return log_dir
def calc_f1(y_true, y_pred, multilabel):
if multilabel:
y_pred[y_pred > 0] = 1
y_pred[y_pred <= 0] = 0
else:
y_pred = np.argmax(y_pred, axis=1)
return f1_score(y_true, y_pred, average="micro"), f1_score(
y_true, y_pred, average="macro"
)
def evaluate(model, g, labels, mask, multilabel=False):
model.eval()
with torch.no_grad():
logits = model(g)
logits = logits[mask]
labels = labels[mask]
f1_mic, f1_mac = calc_f1(
labels.cpu().numpy(), logits.cpu().numpy(), multilabel
)
return f1_mic, f1_mac
# load data of GraphSAINT and convert them to the format of dgl
def load_data(args, multilabel):
if not os.path.exists("graphsaintdata") and not os.path.exists("data"):
raise ValueError("The directory graphsaintdata does not exist!")
elif os.path.exists("graphsaintdata") and not os.path.exists("data"):
os.rename("graphsaintdata", "data")
prefix = "data/{}".format(args.dataset)
DataType = namedtuple("Dataset", ["num_classes", "train_nid", "g"])
adj_full = scipy.sparse.load_npz("./{}/adj_full.npz".format(prefix)).astype(
np.bool_
)
g = dgl.from_scipy(adj_full)
num_nodes = g.num_nodes()
adj_train = scipy.sparse.load_npz(
"./{}/adj_train.npz".format(prefix)
).astype(np.bool_)
train_nid = np.array(list(set(adj_train.nonzero()[0])))
role = json.load(open("./{}/role.json".format(prefix)))
mask = np.zeros((num_nodes,), dtype=bool)
train_mask = mask.copy()
train_mask[role["tr"]] = True
val_mask = mask.copy()
val_mask[role["va"]] = True
test_mask = mask.copy()
test_mask[role["te"]] = True
feats = np.load("./{}/feats.npy".format(prefix))
scaler = StandardScaler()
scaler.fit(feats[train_nid])
feats = scaler.transform(feats)
class_map = json.load(open("./{}/class_map.json".format(prefix)))
class_map = {int(k): v for k, v in class_map.items()}
if multilabel:
# Multi-label binary classification
num_classes = len(list(class_map.values())[0])
class_arr = np.zeros((num_nodes, num_classes))
for k, v in class_map.items():
class_arr[k] = v
else:
num_classes = max(class_map.values()) - min(class_map.values()) + 1
class_arr = np.zeros((num_nodes,))
for k, v in class_map.items():
class_arr[k] = v
g.ndata["feat"] = torch.tensor(feats, dtype=torch.float)
g.ndata["label"] = torch.tensor(
class_arr, dtype=torch.float if multilabel else torch.long
)
g.ndata["train_mask"] = torch.tensor(train_mask, dtype=torch.bool)
g.ndata["val_mask"] = torch.tensor(val_mask, dtype=torch.bool)
g.ndata["test_mask"] = torch.tensor(test_mask, dtype=torch.bool)
data = DataType(g=g, num_classes=num_classes, train_nid=train_nid)
return data