156 lines
5.4 KiB
Python
156 lines
5.4 KiB
Python
import pytorch_lightning as pl
|
|
import torch
|
|
import torchmetrics
|
|
from prototorch.components.components import Components
|
|
from prototorch.functions.distances import euclidean_distance
|
|
from prototorch.functions.similarities import cosine_similarity
|
|
|
|
from .abstract import AbstractPrototypeModel, PrototypeImageModel
|
|
from .glvq import SiameseGLVQ
|
|
|
|
|
|
def rescaled_cosine_similarity(x, y):
|
|
"""Cosine Similarity rescaled to [0, 1]."""
|
|
similarities = cosine_similarity(x, y)
|
|
return (similarities + 1.0) / 2.0
|
|
|
|
|
|
def shift_activation(x):
|
|
return (x + 1.0) / 2.0
|
|
|
|
|
|
def euclidean_similarity(x, y, beta=3):
|
|
d = euclidean_distance(x, y)
|
|
return torch.exp(-d * beta)
|
|
|
|
|
|
class CosineSimilarity(torch.nn.Module):
|
|
def __init__(self, activation=shift_activation):
|
|
super().__init__()
|
|
self.activation = activation
|
|
|
|
def forward(self, x, y):
|
|
epsilon = torch.finfo(x.dtype).eps
|
|
normed_x = (x / x.pow(2).sum(dim=tuple(range(
|
|
1, x.ndim)), keepdim=True).clamp(min=epsilon).sqrt()).flatten(
|
|
start_dim=1)
|
|
normed_y = (y / y.pow(2).sum(dim=tuple(range(
|
|
1, y.ndim)), keepdim=True).clamp(min=epsilon).sqrt()).flatten(
|
|
start_dim=1)
|
|
# normed_x = (x / torch.linalg.norm(x, dim=1))
|
|
diss = torch.inner(normed_x, normed_y)
|
|
return self.activation(diss)
|
|
|
|
|
|
class MarginLoss(torch.nn.modules.loss._Loss):
|
|
def __init__(self,
|
|
margin=0.3,
|
|
size_average=None,
|
|
reduce=None,
|
|
reduction="mean"):
|
|
super().__init__(size_average, reduce, reduction)
|
|
self.margin = margin
|
|
|
|
def forward(self, input_, target):
|
|
dp = torch.sum(target * input_, dim=-1)
|
|
dm = torch.max(input_ - target, dim=-1).values
|
|
return torch.nn.functional.relu(dm - dp + self.margin)
|
|
|
|
|
|
class ReasoningLayer(torch.nn.Module):
|
|
def __init__(self, num_components, num_classes, num_replicas=1):
|
|
super().__init__()
|
|
self.num_replicas = num_replicas
|
|
self.num_classes = num_classes
|
|
probabilities_init = torch.zeros(2, 1, num_components,
|
|
self.num_classes)
|
|
probabilities_init.uniform_(0.4, 0.6)
|
|
self.reasoning_probabilities = torch.nn.Parameter(probabilities_init)
|
|
|
|
@property
|
|
def reasonings(self):
|
|
pk = self.reasoning_probabilities[0]
|
|
nk = (1 - pk) * self.reasoning_probabilities[1]
|
|
ik = 1 - pk - nk
|
|
img = torch.cat([pk, nk, ik], dim=0).permute(1, 0, 2)
|
|
return img.unsqueeze(1)
|
|
|
|
def forward(self, detections):
|
|
pk = self.reasoning_probabilities[0].clamp(0, 1)
|
|
nk = (1 - pk) * self.reasoning_probabilities[1].clamp(0, 1)
|
|
numerator = (detections @ (pk - nk)) + nk.sum(1)
|
|
probs = numerator / (pk + nk).sum(1)
|
|
probs = probs.squeeze(0)
|
|
return probs
|
|
|
|
|
|
class CBC(SiameseGLVQ):
|
|
"""Classification-By-Components."""
|
|
def __init__(self,
|
|
hparams,
|
|
margin=0.1,
|
|
similarity=euclidean_similarity,
|
|
**kwargs):
|
|
super().__init__(hparams, **kwargs)
|
|
self.margin = margin
|
|
self.similarity_fn = kwargs.get("similarity_fn", euclidean_similarity)
|
|
num_components = self.components.shape[0]
|
|
self.reasoning_layer = ReasoningLayer(num_components=num_components,
|
|
num_classes=self.num_classes)
|
|
self.component_layer = self.proto_layer
|
|
|
|
@property
|
|
def components(self):
|
|
return self.prototypes
|
|
|
|
@property
|
|
def reasonings(self):
|
|
return self.reasoning_layer.reasonings.cpu()
|
|
|
|
def forward(self, x):
|
|
components, _ = self.component_layer()
|
|
latent_x = self.backbone(x)
|
|
self.backbone.requires_grad_(self.both_path_gradients)
|
|
latent_components = self.backbone(components)
|
|
self.backbone.requires_grad_(True)
|
|
detections = self.similarity_fn(latent_x, latent_components)
|
|
probs = self.reasoning_layer(detections)
|
|
return probs
|
|
|
|
def shared_step(self, batch, batch_idx, optimizer_idx=None):
|
|
x, y = batch
|
|
# x = x.view(x.size(0), -1)
|
|
y_pred = self(x)
|
|
num_classes = self.reasoning_layer.num_classes
|
|
y_true = torch.nn.functional.one_hot(y.long(), num_classes=num_classes)
|
|
loss = MarginLoss(self.margin)(y_pred, y_true).mean(dim=0)
|
|
return y_pred, loss
|
|
|
|
def training_step(self, batch, batch_idx, optimizer_idx=None):
|
|
y_pred, train_loss = self.shared_step(batch, batch_idx, optimizer_idx)
|
|
preds = torch.argmax(y_pred, dim=1)
|
|
accuracy = torchmetrics.functional.accuracy(preds.int(),
|
|
batch[1].int())
|
|
self.log("train_acc",
|
|
accuracy,
|
|
on_step=False,
|
|
on_epoch=True,
|
|
prog_bar=True,
|
|
logger=True)
|
|
return train_loss
|
|
|
|
def predict(self, x):
|
|
with torch.no_grad():
|
|
y_pred = self(x)
|
|
y_pred = torch.argmax(y_pred, dim=1)
|
|
return y_pred
|
|
|
|
|
|
class ImageCBC(CBC):
|
|
"""CBC model that constrains the components to the range [0, 1] by
|
|
clamping after updates.
|
|
"""
|
|
def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
|
|
# super().on_train_batch_end(outputs, batch, batch_idx, dataloader_idx)
|
|
self.component_layer.components.data.clamp_(0.0, 1.0)
|