prototorch_models/prototorch/models/cbc.py
2021-06-04 22:20:32 +02:00

86 lines
2.8 KiB
Python

import torch
import torchmetrics
from .abstract import ImagePrototypesMixin
from .extras import (
CosineSimilarity,
MarginLoss,
ReasoningLayer,
euclidean_similarity,
rescaled_cosine_similarity,
shift_activation,
)
from .glvq import SiameseGLVQ
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(ImagePrototypesMixin, CBC):
"""CBC model that constrains the components to the range [0, 1] by
clamping after updates.
"""
def __init__(self, hparams, **kwargs):
super().__init__(hparams, **kwargs)
# Namespace hook
self.proto_layer = self.component_layer