prototorch_models/prototorch/models/knn.py

63 lines
1.8 KiB
Python
Raw Normal View History

2021-05-11 15:22:02 +00:00
"""The popular K-Nearest-Neighbors classification algorithm."""
import warnings
import torch
import torchmetrics
from prototorch.components import LabeledComponents
2021-05-17 14:59:35 +00:00
from prototorch.components.initializers import parse_data_arg
2021-05-11 15:22:02 +00:00
from prototorch.functions.competitions import knnc
from prototorch.functions.distances import euclidean_distance
from .abstract import AbstractPrototypeModel
class KNN(AbstractPrototypeModel):
"""K-Nearest-Neighbors classification algorithm."""
def __init__(self, hparams, **kwargs):
super().__init__()
self.save_hyperparameters(hparams)
# Default Values
self.hparams.setdefault("k", 1)
self.hparams.setdefault("distance", euclidean_distance)
data = kwargs.get("data")
2021-05-17 14:59:35 +00:00
x_train, y_train = parse_data_arg(data)
2021-05-11 15:22:02 +00:00
self.proto_layer = LabeledComponents(initialized_components=(x_train,
y_train))
self.train_acc = torchmetrics.Accuracy()
@property
def prototype_labels(self):
2021-05-13 13:22:01 +00:00
return self.proto_layer.component_labels.detach()
2021-05-11 15:22:02 +00:00
def forward(self, x):
protos, _ = self.proto_layer()
dis = self.hparams.distance(x, protos)
return dis
def predict(self, x):
# model.eval() # ?!
with torch.no_grad():
d = self(x)
plabels = self.proto_layer.component_labels
y_pred = knnc(d, plabels, k=self.hparams.k)
2021-05-13 13:22:01 +00:00
return y_pred
2021-05-11 15:22:02 +00:00
def training_step(self, train_batch, batch_idx, optimizer_idx=None):
return 1
def on_train_batch_start(self,
train_batch,
batch_idx,
dataloader_idx=None):
warnings.warn("k-NN has no training, skipping!")
return -1
def configure_optimizers(self):
return None