class RandomWalk():
def __init__(self,
edge_index,
walk_length,
context_size,
walks_per_node=1,
p=1,
q=1,
num_negative_samples=1,
num_nodes=None,
sparse=False):
if random_walk is None:
raise ImportError('`Node2Vec` requires `torch-cluster`.')
N = maybe_num_nodes(edge_index, num_nodes)
row, col = edge_index
self.adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N))
self.adj = self.adj.to('cpu')
assert walk_length >= context_size
self.walk_length = walk_length - 1
self.context_size = context_size
self.walks_per_node = walks_per_node
self.p = p
self.q = q
self.num_negative_samples = num_negative_samples
def loader(self, **kwargs):
return DataLoader(range(self.adj.sparse_size(0)),
collate_fn=self.sample,
**kwargs)
def sample(self, batch):
if not isinstance(batch, torch.Tensor):
batch = torch.tensor(batch)
batch = batch.repeat(self.walks_per_node)
rowptr, col, _ = self.adj.csr()
rw = random_walk(rowptr, col, batch, self.walk_length, self.p, self.q)
if not isinstance(rw, torch.Tensor):
rw = rw[0]
walks = []
num_walks_per_rw = 1 + self.walk_length + 1 - self.context_size
for j in range(num_walks_per_rw):
for i in range(1, self.context_size):
walks.append(rw[:, [j, j + i]])
return torch.cat(walks, dim=0)
class NeighborSampler(torch.utils.data.DataLoader):
r"""The neighbor sampler from the `"Inductive Representation Learning on
Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper, which allows
for mini-batch training of GNNs on large-scale graphs where full-batch
training is not feasible.
Given a GNN with :math:`L` layers and a specific mini-batch of nodes
:obj:`node_idx` for which we want to compute embeddings, this module
iteratively samples neighbors and constructs bipartite graphs that simulate
the actual computation flow of GNNs.
More specifically, :obj:`sizes` denotes how much neighbors we want to
sample for each node in each layer.
This module then takes in these :obj:`sizes` and iteratively samples
:obj:`sizes[l]` for each node involved in layer :obj:`l`.
In the next layer, sampling is repeated for the union of nodes that were
already encountered.
The actual computation graphs are then returned in reverse-mode, meaning
that we pass messages from a larger set of nodes to a smaller one, until we
reach the nodes for which we originally wanted to compute embeddings.
Hence, an item returned by :class:`NeighborSampler` holds the current
:obj:`batch_size`, the IDs :obj:`n_id` of all nodes involved in the
computation, and a list of bipartite graph objects via the tuple
:obj:`(edge_index, e_id, size)`, where :obj:`edge_index` represents the
bipartite edges between source and target nodes, :obj:`e_id` denotes the
IDs of original edges in the full graph, and :obj:`size` holds the shape
of the bipartite graph.
For each bipartite graph, target nodes are also included at the beginning
of the list of source nodes so that one can easily apply skip-connections
or add self-loops.
.. warning::
:class:`~torch_geometric.loader.NeighborSampler` is deprecated and will
be removed in a future release.
Use :class:`torch_geometric.loader.NeighborLoader` instead.
.. note::
For an example of using :obj:`NeighborSampler`, see
`examples/reddit.py
<https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
reddit.py>`_ or
`examples/ogbn_products_sage.py
<https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
ogbn_products_sage.py>`_.
Args:
edge_index (Tensor or SparseTensor): A :obj:`torch.LongTensor` or a
:obj:`torch_sparse.SparseTensor` that defines the underlying graph
connectivity/message passing flow.
:obj:`edge_index` holds the indices of a (sparse) symmetric
adjacency matrix.
If :obj:`edge_index` is of type :obj:`torch.LongTensor`, its shape
must be defined as :obj:`[2, num_edges]`, where messages from nodes
:obj:`edge_index[0]` are sent to nodes in :obj:`edge_index[1]`
(in case :obj:`flow="source_to_target"`).
If :obj:`edge_index` is of type :obj:`torch_sparse.SparseTensor`,
its sparse indices :obj:`(row, col)` should relate to
:obj:`row = edge_index[1]` and :obj:`col = edge_index[0]`.
The major difference between both formats is that we need to input
the *transposed* sparse adjacency matrix.
sizes ([int]): The number of neighbors to sample for each node in each
layer. If set to :obj:`sizes[l] = -1`, all neighbors are included
in layer :obj:`l`.
node_idx (LongTensor, optional): The nodes that should be considered
for creating mini-batches. If set to :obj:`None`, all nodes will be
considered.
num_nodes (int, optional): The number of nodes in the graph.
(default: :obj:`None`)
return_e_id (bool, optional): If set to :obj:`False`, will not return
original edge indices of sampled edges. This is only useful in case
when operating on graphs without edge features to save memory.
(default: :obj:`True`)
transform (callable, optional): A function/transform that takes in
a sampled mini-batch and returns a transformed version.
(default: :obj:`None`)
**kwargs (optional): Additional arguments of
:class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,
:obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.
"""
def __init__(self,
edge_index: Union[Tensor, SparseTensor],
sizes: List[int],
node_idx: Optional[Tensor] = None,
num_nodes: Optional[int] = None,
return_e_id: bool = True,
transform: Callable = None,
**kwargs):
edge_index = edge_index.to('cpu')
if 'collate_fn' in kwargs:
del kwargs['collate_fn']
if 'dataset' in kwargs:
del kwargs['dataset']
# Save for Pytorch Lightning...
self.edge_index = edge_index
self.node_idx = node_idx
self.num_nodes = num_nodes
self.sizes = sizes
self.return_e_id = return_e_id
self.transform = transform
self.is_sparse_tensor = isinstance(edge_index, SparseTensor)
self.__val__ = None
# Obtain a *transposed* `SparseTensor` instance.
if not self.is_sparse_tensor:
if (num_nodes is None and node_idx is not None
and node_idx.dtype == torch.bool):
num_nodes = node_idx.size(0)
if (num_nodes is None and node_idx is not None
and node_idx.dtype == torch.long):
num_nodes = max(int(edge_index.max()), int(node_idx.max())) + 1
if num_nodes is None:
num_nodes = int(edge_index.max()) + 1
value = torch.arange(edge_index.size(1)) if return_e_id else None
self.adj_t = SparseTensor(row=edge_index[0],
col=edge_index[1],
value=value,
sparse_sizes=(num_nodes, num_nodes)).t()
else:
adj_t = edge_index
if return_e_id:
self.__val__ = adj_t.storage.value()
value = torch.arange(adj_t.nnz())
adj_t = adj_t.set_value(value, layout='coo')
self.adj_t = adj_t
self.adj_t.storage.rowptr()
if node_idx is None:
node_idx = torch.arange(self.adj_t.sparse_size(0))
elif node_idx.dtype == torch.bool:
node_idx = node_idx.nonzero(as_tuple=False).view(-1)
super().__init__(node_idx.view(-1).tolist(),
collate_fn=self.sample,
**kwargs)
def sample(self, batch):
if not isinstance(batch, Tensor):
batch = torch.tensor(batch)
batch_size: int = len(batch)
adjs = []
n_id = batch
for size in self.sizes:
adj_t, n_id = self.adj_t.sample_adj(n_id, size, replace=False)
e_id = adj_t.storage.value()
size = adj_t.sparse_sizes()[::-1]
if self.__val__ is not None:
adj_t.set_value_(self.__val__[e_id], layout='coo')
if self.is_sparse_tensor:
adjs.append(Adj(adj_t, e_id, size))
else:
row, col, _ = adj_t.coo()
edge_index = torch.stack([col, row], dim=0)
adjs.append(EdgeIndex(edge_index, e_id, size))
adjs = adjs[0] if len(adjs) == 1 else adjs[::-1]
out = (batch_size, n_id, adjs)
out = self.transform(*out) if self.transform is not None else out
return out
def __repr__(self) -> str:
return f'{self.__class__.__name__}(sizes={self.sizes})'
class Node2Vec(torch.nn.Module):
r"""The Node2Vec model from the
`"node2vec: Scalable Feature Learning for Networks"
<https://arxiv.org/abs/1607.00653>`_ paper where random walks of
length :obj:`walk_length` are sampled in a given graph, and node embeddings
are learned via negative sampling optimization.
.. note::
For an example of using Node2Vec, see `examples/node2vec.py
<https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
node2vec.py>`_.
Args:
edge_index (LongTensor): The edge indices.
embedding_dim (int): The size of each embedding vector.
walk_length (int): The walk length.
context_size (int): The actual context size which is considered for
positive samples. This parameter increases the effective sampling
rate by reusing samples across different source nodes.
walks_per_node (int, optional): The number of walks to sample for each
node. (default: :obj:`1`)
p (float, optional): Likelihood of immediately revisiting a node in the
walk. (default: :obj:`1`)
q (float, optional): Control parameter to interpolate between
breadth-first strategy and depth-first strategy (default: :obj:`1`)
num_negative_samples (int, optional): The number of negative samples to
use for each positive sample. (default: :obj:`1`)
num_nodes (int, optional): The number of nodes. (default: :obj:`None`)
sparse (bool, optional): If set to :obj:`True`, gradients w.r.t. to the
weight matrix will be sparse. (default: :obj:`False`)
"""
def __init__(self, edge_index, embedding_dim, walk_length, context_size,
walks_per_node=1, p=1, q=1, num_negative_samples=1,
num_nodes=None, sparse=False):
super().__init__()
if random_walk is None:
raise ImportError('`Node2Vec` requires `torch-cluster`.')
N = maybe_num_nodes(edge_index, num_nodes)
row, col = edge_index
self.adj = SparseTensor(row=row, col=col, sparse_sizes=(N, N))
self.adj = self.adj.to('cpu')
assert walk_length >= context_size
self.embedding_dim = embedding_dim
self.walk_length = walk_length - 1
self.context_size = context_size
self.walks_per_node = walks_per_node
self.p = p
self.q = q
self.num_negative_samples = num_negative_samples
self.embedding = Embedding(N, embedding_dim, sparse=sparse)
self.reset_parameters()
def reset_parameters(self):
self.embedding.reset_parameters()
def forward(self, batch=None):
"""Returns the embeddings for the nodes in :obj:`batch`."""
emb = self.embedding.weight
return emb if batch is None else emb.index_select(0, batch)
def loader(self, **kwargs):
return DataLoader(range(self.adj.sparse_size(0)),
collate_fn=self.sample, **kwargs)
def pos_sample(self, batch):
batch = batch.repeat(self.walks_per_node)
rowptr, col, _ = self.adj.csr()
rw = random_walk(rowptr, col, batch, self.walk_length, self.p, self.q)
if not isinstance(rw, torch.Tensor):
rw = rw[0]
walks = []
num_walks_per_rw = 1 + self.walk_length + 1 - self.context_size
for j in range(num_walks_per_rw):
walks.append(rw[:, j:j + self.context_size])
return torch.cat(walks, dim=0)
def neg_sample(self, batch):
batch = batch.repeat(self.walks_per_node * self.num_negative_samples)
rw = torch.randint(self.adj.sparse_size(0),
(batch.size(0), self.walk_length))
rw = torch.cat([batch.view(-1, 1), rw], dim=-1)
walks = []
num_walks_per_rw = 1 + self.walk_length + 1 - self.context_size
for j in range(num_walks_per_rw):
walks.append(rw[:, j:j + self.context_size])
return torch.cat(walks, dim=0)
def sample(self, batch):
if not isinstance(batch, torch.Tensor):
batch = torch.tensor(batch)
return self.pos_sample(batch), self.neg_sample(batch)
def loss(self, pos_rw, neg_rw):
r"""Computes the loss given positive and negative random walks."""
# Positive loss.
start, rest = pos_rw[:, 0], pos_rw[:, 1:].contiguous()
h_start = self.embedding(start).view(pos_rw.size(0), 1,
self.embedding_dim)
h_rest = self.embedding(rest.view(-1)).view(pos_rw.size(0), -1,
self.embedding_dim)
out = (h_start * h_rest).sum(dim=-1).view(-1)
pos_loss = -torch.log(torch.sigmoid(out) + EPS).mean()
# Negative loss.
start, rest = neg_rw[:, 0], neg_rw[:, 1:].contiguous()
h_start = self.embedding(start).view(neg_rw.size(0), 1,
self.embedding_dim)
h_rest = self.embedding(rest.view(-1)).view(neg_rw.size(0), -1,
self.embedding_dim)
out = (h_start * h_rest).sum(dim=-1).view(-1)
neg_loss = -torch.log(1 - torch.sigmoid(out) + EPS).mean()
return pos_loss + neg_loss
def test(self, train_z, train_y, test_z, test_y, solver='lbfgs',
multi_class='auto', *args, **kwargs):
r"""Evaluates latent space quality via a logistic regression downstream
task."""
from sklearn.linear_model import LogisticRegression
clf = LogisticRegression(solver=solver, multi_class=multi_class, *args,
**kwargs).fit(train_z.detach().cpu().numpy(),
train_y.detach().cpu().numpy())
return clf.score(test_z.detach().cpu().numpy(),
test_y.detach().cpu().numpy())
def __repr__(self) -> str:
return (f'{self.__class__.__name__}({self.embedding.weight.size(0)}, '
f'{self.embedding.weight.size(1)})')
class HeteroNeighborSampler(torch.utils.data.DataLoader):
def __init__(self,
edge_index: Union[Tensor, SparseTensor],
sizes: List[int],
node_idx: Optional[Tensor] = None,
num_nodes: Optional[int] = None,
return_e_id: bool = True,
**kwargs):
"""
Args:
edge_index:
sizes:
node_idx:
num_nodes:
return_e_id (bool):
**kwargs:
"""
self.sizes = sizes
self.return_e_id = return_e_id
self.is_sparse_tensor = isinstance(edge_index, SparseTensor)
self.__val__ = None
# Obtain a *transposed* `SparseTensor` instance.
edge_index = edge_index.to('cpu')
if not self.is_sparse_tensor:
num_nodes = maybe_num_nodes(edge_index, num_nodes)
value = torch.arange(edge_index.size(1)) if return_e_id else None
# Sampling source_to_target
self.adj_t = SparseTensor(row=edge_index[1],
col=edge_index[0],
value=value,
sparse_sizes=(num_nodes, num_nodes)).t()
else:
adj_t = edge_index
if return_e_id:
self.__val__ = adj_t.storage.value()
value = torch.arange(adj_t.nnz())
adj_t = adj_t.set_value(value, layout='coo')
self.adj_t = adj_t
self.adj_t.storage.rowptr()
if node_idx is None:
node_idx = torch.arange(self.adj_t.sparse_size(0))
elif node_idx.dtype == torch.bool:
node_idx = node_idx.nonzero(as_tuple=False).view(-1)
super(HeteroNeighborSampler, self).__init__(node_idx.view(-1).tolist(),
collate_fn=self.sample,
**kwargs)
def sample(self, batch):
"""
Args:
batch:
"""
if not isinstance(batch, Tensor):
batch = torch.tensor(batch)
batch_size: int = len(batch)
adjs = []
n_id = batch
for size in self.sizes:
adj_t, n_id = self.adj_t.sample_adj(n_id, size, replace=False)
e_id = adj_t.storage.value()
size = adj_t.sparse_sizes()[::-1]
if self.__val__ is not None:
adj_t.set_value_(self.__val__[e_id], layout='coo')
if self.is_sparse_tensor:
adjs.append(Adj(adj_t, e_id, size))
else:
row, col, _ = adj_t.coo()
edge_index = torch.stack([row, col], dim=0)
adjs.append(EdgeIndex(edge_index, e_id, size))
return batch_size, n_id, adjs
class ShaDowKHopSampler(torch.utils.data.DataLoader):
r"""The ShaDow :math:`k`-hop sampler from the `"Deep Graph Neural Networks
with Shallow Subgraph Samplers" <https://arxiv.org/abs/2012.01380>`_ paper.
Given a graph in a :obj:`data` object, the sampler will create shallow,
localized subgraphs.
A deep GNN on this local graph then smooths the informative local signals.
Args:
data (torch_geometric.data.Data): The graph data object.
depth (int): The depth/number of hops of the localized subgraph.
num_neighbors (int): The number of neighbors to sample for each node in
each hop.
node_idx (LongTensor or BoolTensor, optional): The nodes that should be
considered for creating mini-batches.
If set to :obj:`None`, all nodes will be
considered.
replace (bool, optional): If set to :obj:`True`, will sample neighbors
with replacement. (default: :obj:`False`)
**kwargs (optional): Additional arguments of
:class:`torch.utils.data.DataLoader`, such as :obj:`batch_size` or
:obj:`num_workers`.
"""
def __init__(self,
data: Data,
depth: int,
num_neighbors: int,
node_idx: Optional[Tensor] = None,
replace: bool = False,
**kwargs):
self.data = copy.copy(data)
self.depth = depth
self.num_neighbors = num_neighbors
self.replace = replace
if data.edge_index is not None:
self.is_sparse_tensor = False
row, col = data.edge_index.cpu()
self.adj_t = SparseTensor(row=row,
col=col,
value=torch.arange(col.size(0)),
sparse_sizes=(data.num_nodes,
data.num_nodes)).t()
else:
self.is_sparse_tensor = True
self.adj_t = data.adj_t.cpu()
if node_idx is None:
node_idx = torch.arange(self.adj_t.sparse_size(0))
elif node_idx.dtype == torch.bool:
node_idx = node_idx.nonzero(as_tuple=False).view(-1)
self.node_idx = node_idx
super().__init__(node_idx.tolist(),
collate_fn=self.__collate__,
**kwargs)
def __collate__(self, n_id):
n_id = torch.tensor(n_id)
rowptr, col, value = self.adj_t.csr()
out = torch.ops.torch_sparse.ego_k_hop_sample_adj(
rowptr, col, n_id, self.depth, self.num_neighbors, self.replace)
rowptr, col, n_id, e_id, ptr, root_n_id = out
adj_t = SparseTensor(rowptr=rowptr,
col=col,
value=value[e_id] if value is not None else None,
sparse_sizes=(n_id.numel(), n_id.numel()),
is_sorted=True)
batch = Batch(batch=torch.ops.torch_sparse.ptr2ind(ptr, n_id.numel()),
ptr=ptr)
batch.root_n_id = root_n_id
if self.is_sparse_tensor:
batch.adj_t = adj_t
else:
row, col, e_id = adj_t.t().coo()
batch.edge_index = torch.stack([row, col], dim=0)
for k, v in self.data:
if k in ['edge_index', 'adj_t', 'num_nodes']:
continue
if k == 'y' and v.size(0) == self.data.num_nodes:
batch[k] = v[n_id][root_n_id]
elif isinstance(v, Tensor) and v.size(0) == self.data.num_nodes:
batch[k] = v[n_id]
elif isinstance(v, Tensor) and v.size(0) == self.data.num_edges:
batch[k] = v[e_id]
else:
batch[k] = v
return batch
