FAISSKNN

faissknn 包含了多类和多标签K-最近邻分类器的实现。分类器遵循 scikit-learn 的 fit、predict 和 predict_proba 方法。

安装

FAISS的作者建议通过conda安装 faiss，例如 conda install -c pytorch faiss-gpu。有关更多信息，请参阅 FAISS 安装页面。

安装 faiss 后，可以通过PyPI安装 faissknn

pip install faissknn

用法

多类

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)

多标签

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 = FaissKNNClassifier(
    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 还支持在GPU上运行以提高计算速度。只需将设备更改为 cuda 或特定的cuda设备 cuda:0

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