torch_geometric.nn.pool.ApproxMIPSKNNIndex
- class ApproxMIPSKNNIndex(num_cells: int, num_cells_to_visit: int, bits_per_vector: int, emb: Optional[Tensor] = None, reserve: Optional[int] = None)[source]
Bases:
KNNIndex
Performs fast approximate \(k\)-nearest neighbor search (\(k\)-NN) based on the maximum inner product via the
faiss
library. Hyperparameters needs to be tuned for speed-accuracy trade-off.- Parameters:
num_cells (int) – The number of cells.
num_cells_to_visit (int) – The number of cells that are visited to perform to search.
bits_per_vector (int) – The number of bits per sub-vector.
emb (torch.Tensor, optional) – The data points to add. (default:
None
)reserve (int, optional) – The number of elements to reserve memory for before re-allocating (GPU only). (default:
None
)