import torch import torchmetrics from ..core.components import ReasoningComponents 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, **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