torch_geometric.transforms.KNNGraph

class KNNGraph(k: int = 6, loop: bool = False, force_undirected: bool = False, flow: str = 'source_to_target', cosine: bool = False, num_workers: int = 1)[source]

Bases: BaseTransform

Creates a k-NN graph based on node positions data.pos (functional name: knn_graph).

Parameters

  • k (int, optional) – The number of neighbors. (default: 6)

  • loop (bool, optional) – If True, the graph will contain self-loops. (default: False)

  • force_undirected (bool, optional) – If set to True, new edges will be undirected. (default: False)

  • flow (str, optional) – The flow direction when used in combination with message passing ("source_to_target" or "target_to_source"). If set to "source_to_target", every target node will have exactly \(k\) source nodes pointing to it. (default: "source_to_target")

  • cosine (bool, optional) – If True, will use the cosine distance instead of euclidean distance to find nearest neighbors. (default: False)

  • num_workers (int) – Number of workers to use for computation. Has no effect in case batch is not None, or the input lies on the GPU. (default: 1)