prototorch_models/prototorch/models/probabilistic.py

64 lines
2.0 KiB
Python
Raw Normal View History

2021-05-25 18:26:15 +00:00
"""Probabilistic GLVQ methods"""
import torch
2021-05-28 15:13:06 +00:00
from prototorch.functions.competitions import stratified_sum
from prototorch.functions.losses import (log_likelihood_ratio_loss,
robust_soft_loss)
2021-05-28 15:13:06 +00:00
from prototorch.functions.transform import gaussian
2021-05-25 18:26:15 +00:00
from .glvq import GLVQ
2021-05-28 15:13:06 +00:00
class ProbabilisticLVQ(GLVQ):
def __init__(self, hparams, rejection_confidence=1.0, **kwargs):
2021-05-25 18:26:15 +00:00
super().__init__(hparams, **kwargs)
2021-05-28 15:13:06 +00:00
self.conditional_distribution = gaussian
self.rejection_confidence = rejection_confidence
2021-05-25 18:26:15 +00:00
def predict(self, x):
probabilities = self.forward(x)
confidence, prediction = torch.max(probabilities, dim=1)
2021-05-28 15:13:06 +00:00
prediction[confidence < self.rejection_confidence] = -1
2021-05-25 18:26:15 +00:00
return prediction
def forward(self, x):
distances = self._forward(x)
conditional = self.conditional_distribution(distances,
self.hparams.variance)
prior = 1.0 / torch.Tensor(self.proto_layer.distribution).sum().item()
posterior = conditional * prior
plabels = torch.LongTensor(self.proto_layer.component_labels)
2021-05-28 15:13:06 +00:00
y_pred = stratified_sum(posterior.T, plabels)
2021-05-25 18:26:15 +00:00
return y_pred
def training_step(self, batch, batch_idx, optimizer_idx=None):
X, y = batch
out = self.forward(X)
plabels = self.proto_layer.component_labels
2021-05-28 15:13:06 +00:00
batch_loss = -self.loss_fn(out, y, plabels)
2021-05-25 18:26:15 +00:00
loss = batch_loss.sum(dim=0)
return loss
2021-05-28 15:13:06 +00:00
class LikelihoodRatioLVQ(ProbabilisticLVQ):
"""Learning Vector Quantization based on Likelihood Ratios
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.loss_fn = log_likelihood_ratio_loss
2021-05-28 15:13:06 +00:00
class RSLVQ(ProbabilisticLVQ):
"""Learning Vector Quantization based on Likelihood Ratios
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.loss_fn = robust_soft_loss
2021-05-25 18:26:15 +00:00
2021-05-28 15:13:06 +00:00
__all__ = ["LikelihoodRatioLVQ", "RSLVQ"]