using softmax classifier as evaluator

Andrewzh112 · Dec 5, 2020 · 4656d16 · 4656d16
1 parent 5c088b9
commit 4656d16
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/moco/main.py b/moco/main.py
@@ -119,7 +119,7 @@
         feature_bank = torch.cat(feature_bank).cpu().numpy()
         feature_labels = torch.cat(feature_labels).numpy()
 
-        linear_classifier = LogisticRegression()
+        linear_classifier = LogisticRegression(multi_class='multinomial', solver='lbfgs')
         linear_classifier.fit(feature_bank, feature_labels)
 
         y_preds, y_trues = [], []

diff --git a/moco/utils.py b/moco/utils.py
@@ -19,9 +19,9 @@ def forward(self, q, k, memo_bank):
         pos_logits = torch.einsum('ij,ij->i', [q, k]).unsqueeze(-1)
         neg_logits = torch.einsum('ij,kj->ik', [q, memo_bank.queue.clone()])
         logits = torch.cat([pos_logits, neg_logits], dim=1)
-    
+
         # zero is the positive "class"
-        labels = torch.new_zeros(N)
+        labels = logits.new_zeros(N, dtype=torch.long)
         return self.criterion(logits / self.T, labels)
 
 

diff --git a/simsiam/main.py b/simsiam/main.py
@@ -113,7 +113,7 @@
         feature_bank = torch.cat(feature_bank, dim=0).cpu().numpy()
         feature_labels = torch.cat(targets, dim=0).numpy()
 
-        linear_classifier = LogisticRegression()
+        linear_classifier = LogisticRegression(multi_class='multinomial', solver='lbfgs')
         linear_classifier.fit(feature_bank, feature_labels)
 
         y_preds, y_trues = [], []