fixed forwarding of temperature scaling in classifier with embedding

Author: Dingel321

src/cryo_sbi/inference/models/estimator_models.py

@@ -48,23 +48,30 @@ def __init__(
         self.classifier.append(nn.Linear(nodes_per_layer, num_classes))
         self.classifier = nn.Sequential(*self.classifier)
 
-    def forward(self, z):
-        return self.classifier(z)
+    def forward(self, z, tau=1.0):
+        return self.classifier(z) / tau
 
 
 @add_classifier("PROTOTYPE")
 class PrototypeClassifier(BaseClassifier):
-    def __init__(self, input_dim, num_classes):
+    def __init__(self, input_dim, num_classes, noise_scale=0.0):
         super().__init__(input_dim, num_classes)
 
         self.input_dim = input_dim
         self.num_classes = num_classes
+        self.noise_scale = noise_scale
         self.prototypes = nn.Parameter(torch.randn(self.num_classes, self.input_dim))
 
     def forward(self, z, tau=1.0):
+        if self.training and self.noise_scale > 0.0:
+            prototypes = self.prototypes + torch.randn_like(self.prototypes) * self.noise_scale
+        else:
+            prototypes = self.prototypes
+
         z2 = (z**2).sum(dim=1, keepdim=True)
-        p2 = (self.prototypes**2).sum(dim=1).unsqueeze(0)
-        logits = -(z2 + p2 - 2 * z @ self.prototypes.T) / tau
+        p2 = (prototypes**2).sum(dim=1).unsqueeze(0)
+        logits = -(z2 + p2 - 2 * z @ prototypes.T) / tau
+
         return logits
 
 
@@ -78,13 +85,12 @@ def __init__(
         self.classifier = classifier()
         self.embedding = embedding_net()
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.classifier(self.embedding(x))
-
-    def probs(self, x: torch.Tensor) -> torch.Tensor:
-        return torch.nn.functional.softmax(self.forward(x), dim=1)
+    def forward(self, x: torch.Tensor, tau=1.0) -> torch.Tensor:
+        return self.classifier(self.embedding(x), tau=tau)
 
-    def logits_embedding(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def probs(self, x: torch.Tensor, tau=1.0) -> torch.Tensor:
+        return torch.nn.functional.softmax(self.forward(x, tau=tau) / tau, dim=1)
+    def logits_embedding(self, x: torch.Tensor, tau=1.0) -> Tuple[torch.Tensor, torch.Tensor]:
         embeddings = self.embedding(x)
-        logits = self.classifier(embeddings)
+        logits = self.classifier(embeddings, tau=tau)
         return logits, embeddings