prototorch_models/prototorch/models/probabilistic.py

100 lines
3.5 KiB
Python
Raw Normal View History

2021-05-25 18:26:15 +00:00
"""Probabilistic GLVQ methods"""
import torch
2021-06-04 20:20:32 +00:00
from prototorch.functions.losses import nllr_loss, rslvq_loss
2021-06-08 13:01:08 +00:00
from prototorch.functions.pooling import (stratified_min_pooling,
stratified_sum_pooling)
from prototorch.functions.transforms import (GaussianPrior,
RankScaledGaussianPrior)
2021-06-04 20:20:32 +00:00
from prototorch.modules import LambdaLayer, LossLayer
2021-05-25 18:26:15 +00:00
2021-06-08 13:01:08 +00:00
from .glvq import GLVQ, SiameseGMLVQ
2021-05-25 18:26:15 +00:00
2021-06-01 21:39:06 +00:00
class CELVQ(GLVQ):
"""Cross-Entropy Learning Vector Quantization."""
def __init__(self, hparams, **kwargs):
super().__init__(hparams, **kwargs)
2021-06-04 20:20:32 +00:00
# Loss
2021-06-01 21:39:06 +00:00
self.loss = torch.nn.CrossEntropyLoss()
def shared_step(self, batch, batch_idx, optimizer_idx=None):
x, y = batch
2021-06-04 20:20:32 +00:00
out = self.compute_distances(x) # [None, num_protos]
2021-06-01 21:39:06 +00:00
plabels = self.proto_layer.component_labels
2021-06-04 20:20:32 +00:00
winning = stratified_min_pooling(out, plabels) # [None, num_classes]
probs = -1.0 * winning
2021-06-01 21:39:06 +00:00
batch_loss = self.loss(probs, y.long())
loss = batch_loss.sum(dim=0)
return out, loss
2021-05-28 15:13:06 +00:00
class ProbabilisticLVQ(GLVQ):
2021-05-31 15:56:45 +00:00
def __init__(self, hparams, rejection_confidence=0.0, **kwargs):
2021-05-25 18:26:15 +00:00
super().__init__(hparams, **kwargs)
2021-06-08 13:01:08 +00:00
self.conditional_distribution = None
2021-05-28 15:13:06 +00:00
self.rejection_confidence = rejection_confidence
2021-05-25 18:26:15 +00:00
def forward(self, x):
2021-06-04 20:20:32 +00:00
distances = self.compute_distances(x)
2021-06-08 13:01:08 +00:00
conditional = self.conditional_distribution(distances)
prior = (1. / self.num_prototypes) * torch.ones(self.num_prototypes,
device=self.device)
2021-05-25 18:26:15 +00:00
posterior = conditional * prior
2021-06-01 15:44:10 +00:00
plabels = self.proto_layer._labels
2021-06-04 20:20:32 +00:00
y_pred = stratified_sum_pooling(posterior, plabels)
2021-05-25 18:26:15 +00:00
return y_pred
2021-06-01 15:44:10 +00:00
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
2021-05-25 18:26:15 +00:00
def training_step(self, batch, batch_idx, optimizer_idx=None):
2021-06-04 20:20:32 +00:00
x, y = batch
out = self.forward(x)
2021-05-25 18:26:15 +00:00
plabels = self.proto_layer.component_labels
2021-06-04 20:20:32 +00:00
batch_loss = self.loss(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 SLVQ(ProbabilisticLVQ):
"""Soft Learning Vector Quantization."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
2021-06-04 20:20:32 +00:00
self.loss = LossLayer(nllr_loss)
2021-06-08 13:01:08 +00:00
self.conditional_distribution = GaussianPrior(self.hparams.variance)
2021-05-28 15:13:06 +00:00
class RSLVQ(ProbabilisticLVQ):
2021-05-31 15:56:45 +00:00
"""Robust Soft Learning Vector Quantization."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
2021-06-04 20:20:32 +00:00
self.loss = LossLayer(rslvq_loss)
2021-06-08 13:01:08 +00:00
self.conditional_distribution = GaussianPrior(self.hparams.variance)
class PLVQ(ProbabilisticLVQ, SiameseGMLVQ):
"""Probabilistic Learning Vector Quantization.
TODO: Use Backbone LVQ instead
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.conditional_distribution = RankScaledGaussianPrior(
self.hparams.lambd)
self.loss = torch.nn.KLDivLoss()
def training_step(self, batch, batch_idx, optimizer_idx=None):
x, y = batch
out = self.forward(x)
y_dist = torch.nn.functional.one_hot(
y.long(), num_classes=self.num_classes).float()
batch_loss = self.loss(out, y_dist)
loss = batch_loss.sum(dim=0)
return loss