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 notNone
, or the input lies on the GPU. (default:1
)