faissknn 0.2.0

FAISS implementation of multiclass and multilabel K-Nearest Neighbors Classifiers.

faissknn

faissknn contains implementations for both multiclass and multilabel K-Nearest Neighbors Classifier implementations. The classifiers follow the scikit-learn: fit, predict, and predict_proba methods.

Install

pip install faissknn

Pulls in faiss-cpu along with numpy and torch. Works everywhere — no GPU, CUDA driver, or system toolkit required.

GPU acceleration (CUDA 12.x)

For GPU-enabled FAISS on CUDA 12.x hosts (NVIDIA driver R525+), install the [cuda] extra, which adds faiss-cuda-cu128 (Taylor Geospatial's GPU wheels for CUDA 12.8). No system CUDA toolkit needed — the runtime libraries come from nvidia-cuda-runtime-cu12 / nvidia-cublas-cu12 on PyPI.

pip install "faissknn[cuda]"
pip uninstall -y faiss-cpu
pip install --force-reinstall faiss-cuda-cu128

These wheels also run on CUDA 13 hosts (driver R580+) via NVIDIA's forward-compat guarantee — you just don't get the sm_100 (Blackwell) arch.

Blackwell users (B100 / B200)

If you need sm_100 baked in, use the CUDA 13 wheel via the [cu13] extra, which adds faiss-cuda:

pip install "faissknn[cu13]"
pip uninstall -y faiss-cpu
pip install --force-reinstall faiss-cuda

Why the uninstall+reinstall? faiss-cpu, faiss-cuda-cu128, and faiss-cuda all ship the same faiss/ module, so they can't cleanly coexist. The extra adds the GPU package to the resolution, but removing faiss-cpu and force-reinstalling is what actually gives you a clean GPU install. uv/pip can't auto-detect the host CUDA driver, so this is a one-time manual choice.

Usage

Multiclass:

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split

from faissknn import FaissKNNClassifier

x, y = make_classification()
x_train, x_test, y_train, y_test = train_test_split(x, y)
model = FaissKNNClassifier(
    n_neighbors=5,
    n_classes=None,
    device="cpu"
)
model.fit(x_train, y_train)

y_pred = model.predict(x_test) # (N,)
y_proba = model.predict_proba(x_test) # (N, C)

Multilabel:

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split

from faissknn import FaissKNNMultilabelClassifier

x, y = make_multilabel_classification()
x_train, x_test, y_train, y_test = train_test_split(x, y)
model = FaissKNNMultilabelClassifier(
    n_neighbors=5,
    device="cpu"
)
model.fit(x_train, y_train)

y_pred = model.predict(x_test) # (N, C)
y_proba = model.predict_proba(x_test) # (N, C)

GPU/CUDA: faissknn also supports running on the GPU to speed up computation. Simply change the device to cuda or a specific cuda device cuda:0

model = FaissKNNClassifier(
    n_neighbors=5,
    device="cuda"
)
model = FaissKNNClassifier(
    n_neighbors=5,
    device="cuda:0"
)

Cite

If you use faissknn in your research, please considering citing!

@software{isaac_corley_2026_18370748,
  author       = {Isaac Corley},
  title        = {isaaccorley/faissknn: Zenodo Cite},
  month        = jan,
  year         = 2026,
  publisher    = {Zenodo},
  version      = {v0.0.3},
  doi          = {10.5281/zenodo.18370748},
  url          = {https://doi.org/10.5281/zenodo.18370748},
}