chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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()
|
||||
Reference in New Issue
Block a user