Removed iterations and optimized further.

vit_zana_v3.py

@@ -203,6 +203,15 @@ class BaselineViT(nn.Module):
         # Temperature for similarity-based classification
         self.temperature = nn.Parameter(torch.ones(1) * np.log(1/0.07))
         
+        self.register_buffer(
+            'all_centroids',
+            torch.stack([penta.centroid for penta in self.class_pentachora])
+        )
+        self.register_buffer(
+            'all_centroids_norm',
+            F.normalize(self.all_centroids, dim=-1)
+        )
+
         # Initialize weights
         self.init_weights()
     
@@ -220,29 +229,27 @@ class BaselineViT(nn.Module):
                 nn.init.ones_(m.weight)
                 nn.init.zeros_(m.bias)
     
+    # Then get_class_centroids becomes:
     def get_class_centroids(self) -> torch.Tensor:
-        """Get all class centroids for similarity computation."""
-        centroids = torch.stack([
-            penta.get_centroid() for penta in self.class_pentachora
-        ])
-        return F.normalize(centroids, dim=-1)
-    
+        return self.all_centroids_norm
+        
     def compute_pentachora_similarities(self, features: torch.Tensor) -> torch.Tensor:
         """
-        Compute similarities between features and all class pentachora.
-        
-        Args:
-            features: [batch, dim] features to compare
-        
-        Returns:
-            similarities: [batch, num_classes]
+        Compute similarities between features and all class pentachora (vectorized).
         """
-        similarities = []
-        for penta in self.class_pentachora:
-            sim = penta.compute_similarity(features, mode=self.similarity_mode)
-            similarities.append(sim)
-        
-        return torch.stack(similarities, dim=-1)
+        if self.similarity_mode == 'rose':
+            # Stack all vertices into single tensor for batch Rose scoring
+            all_vertices = torch.stack([penta.vertices for penta in self.class_pentachora])  # [100, 5, 128]
+            # Expand features for batch computation
+            features_exp = features.unsqueeze(1).expand(-1, self.num_classes, -1)  # [B, 100, 128]
+            # Compute Rose scores in parallel
+            return PentachoronStabilizer.rose_score_magnitude(features_exp.reshape(-1, 128), all_vertices.repeat(features.shape[0], 1, 1)).reshape(features.shape[0], -1)
+        else:
+            # Stack all centroids
+            centroids = torch.stack([penta.centroid_norm for penta in self.class_pentachora])  # [100, 128]
+            features_norm = F.normalize(features, dim=-1)  # [B, 128]
+            return torch.matmul(features_norm, centroids.T)  # [B, 100]
+
     
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         """Extract features from images."""