torch_geometric.nn.models.Node2Vec

class Node2Vec(edge_index: Tensor, embedding_dim: int, walk_length: int, context_size: int, walks_per_node: int = 1, p: float = 1.0, q: float = 1.0, num_negative_samples: int = 1, num_nodes: Optional[int] = None, sparse: bool = False)[source]

Bases: Module

The Node2Vec model from the “node2vec: Scalable Feature Learning for Networks” paper where random walks of length 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.

Parameters:

  • edge_index (torch.Tensor) – 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: 1)

  • p (float, optional) – Likelihood of immediately revisiting a node in the walk. (default: 1)

  • q (float, optional) – Control parameter to interpolate between breadth-first strategy and depth-first strategy (default: 1)

  • num_negative_samples (int, optional) – The number of negative samples to use for each positive sample. (default: 1)

  • num_nodes (int, optional) – The number of nodes. (default: None)

  • sparse (bool, optional) – If set to True, gradients w.r.t. to the weight matrix will be sparse. (default: False)

forward(batch: Optional[Tensor] = None) → Tensor[source]

Returns the embeddings for the nodes in batch.

reset_parameters()[source]

Resets all learnable parameters of the module.

loss(pos_rw: Tensor, neg_rw: Tensor) → Tensor[source]

Computes the loss given positive and negative random walks.

test(train_z: Tensor, train_y: Tensor, test_z: Tensor, test_y: Tensor, solver: str = 'lbfgs', multi_class: str = 'auto', *args, **kwargs) → float[source]

Evaluates latent space quality via a logistic regression downstream task.