132 lines
4.0 KiB
Python
132 lines
4.0 KiB
Python
"""Probabilistic GLVQ methods"""
|
|
|
|
import torch
|
|
from prototorch.core.losses import nllr_loss, rslvq_loss
|
|
from prototorch.core.pooling import (
|
|
stratified_min_pooling,
|
|
stratified_sum_pooling,
|
|
)
|
|
from prototorch.nn.wrappers import LossLayer
|
|
|
|
from .extras import GaussianPrior, RankScaledGaussianPrior
|
|
from .glvq import GLVQ, SiameseGMLVQ
|
|
|
|
|
|
class CELVQ(GLVQ):
|
|
"""Cross-Entropy Learning Vector Quantization."""
|
|
|
|
def __init__(self, hparams, **kwargs):
|
|
super().__init__(hparams, **kwargs)
|
|
|
|
# Loss
|
|
self.loss = torch.nn.CrossEntropyLoss()
|
|
|
|
def shared_step(self, batch, batch_idx, optimizer_idx=None):
|
|
x, y = batch
|
|
out = self.compute_distances(x) # [None, num_protos]
|
|
_, plabels = self.proto_layer()
|
|
winning = stratified_min_pooling(out, plabels) # [None, num_classes]
|
|
probs = -1.0 * winning
|
|
batch_loss = self.loss(probs, y.long())
|
|
loss = batch_loss.sum()
|
|
return out, loss
|
|
|
|
|
|
class ProbabilisticLVQ(GLVQ):
|
|
|
|
def __init__(self, hparams, rejection_confidence=0.0, **kwargs):
|
|
super().__init__(hparams, **kwargs)
|
|
|
|
self.rejection_confidence = rejection_confidence
|
|
self._conditional_distribution = None
|
|
|
|
def forward(self, x):
|
|
distances = self.compute_distances(x)
|
|
|
|
conditional = self.conditional_distribution(distances)
|
|
prior = (1. / self.num_prototypes) * torch.ones(self.num_prototypes,
|
|
device=self.device)
|
|
posterior = conditional * prior
|
|
|
|
plabels = self.proto_layer._labels
|
|
if isinstance(plabels, torch.LongTensor) or isinstance(
|
|
plabels, torch.cuda.LongTensor): # type: ignore
|
|
y_pred = stratified_sum_pooling(posterior, plabels) # type: ignore
|
|
else:
|
|
raise ValueError("Labels must be LongTensor.")
|
|
|
|
return y_pred
|
|
|
|
def predict(self, x):
|
|
y_pred = self.forward(x)
|
|
confidence, prediction = torch.max(y_pred, dim=1)
|
|
prediction[confidence < self.rejection_confidence] = -1
|
|
return prediction
|
|
|
|
def training_step(self, batch, batch_idx, optimizer_idx=None):
|
|
x, y = batch
|
|
out = self.forward(x)
|
|
_, plabels = self.proto_layer()
|
|
batch_loss = self.loss(out, y, plabels)
|
|
loss = batch_loss.sum()
|
|
return loss
|
|
|
|
def conditional_distribution(self, distances):
|
|
"""Conditional distribution of distances."""
|
|
if self._conditional_distribution is None:
|
|
raise ValueError("Conditional distribution is not set.")
|
|
return self._conditional_distribution(distances)
|
|
|
|
|
|
class SLVQ(ProbabilisticLVQ):
|
|
"""Soft Learning Vector Quantization."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
# Default hparams
|
|
self.hparams.setdefault("variance", 1.0)
|
|
variance = self.hparams.get("variance")
|
|
|
|
self._conditional_distribution = GaussianPrior(variance)
|
|
self.loss = LossLayer(nllr_loss)
|
|
|
|
|
|
class RSLVQ(ProbabilisticLVQ):
|
|
"""Robust Soft Learning Vector Quantization."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
# Default hparams
|
|
self.hparams.setdefault("variance", 1.0)
|
|
variance = self.hparams.get("variance")
|
|
|
|
self._conditional_distribution = GaussianPrior(variance)
|
|
self.loss = LossLayer(rslvq_loss)
|
|
|
|
|
|
class PLVQ(ProbabilisticLVQ, SiameseGMLVQ):
|
|
"""Probabilistic Learning Vector Quantization.
|
|
|
|
TODO: Use Backbone LVQ instead
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
# Default hparams
|
|
self.hparams.setdefault("lambda", 1.0)
|
|
lam = self.hparams.get("lambda", 1.0)
|
|
|
|
self.conditional_distribution = RankScaledGaussianPrior(lam)
|
|
self.loss = torch.nn.KLDivLoss()
|
|
|
|
# FIXME
|
|
# def training_step(self, batch, batch_idx, optimizer_idx=None):
|
|
# x, y = batch
|
|
# y_pred = self(x)
|
|
# batch_loss = self.loss(y_pred, y)
|
|
# loss = batch_loss.sum()
|
|
# return loss
|