Fix checkpoint loading for KNN eval (#376)

Author: vturrisi

main_knn.py

@@ -24,6 +24,7 @@
 
 import torch
 import torch.nn as nn
+from omegaconf import OmegaConf
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
@@ -126,12 +127,10 @@ def main():
     # load arguments
     with open(args_path) as f:
         method_args = json.load(f)
+    cfg = OmegaConf.create(method_args)
 
     # build the model
-    model = METHODS[method_args["method"]].load_from_checkpoint(
-        ckpt_path, strict=False, **method_args
-    )
-    model.cuda()
+    model = METHODS[method_args["method"]].load_from_checkpoint(ckpt_path, strict=False, cfg=cfg)
 
     # prepare data
     _, T = prepare_transforms(args.dataset)

solo/methods/mocov3.py

@@ -171,8 +171,9 @@ def forward(self, X: torch.Tensor) -> Dict[str, Any]:
         """
 
         out = super().forward(X)
-        q = self.predictor(self.projector(out["feats"]))
-        out.update({"q": q})
+        z = self.projector(out["feats"])
+        q = self.predictor(z)
+        out.update({"q": q, "z": z})
         return out
 
     @torch.no_grad()

solo/methods/swav.py

@@ -43,7 +43,8 @@ def __init__(self, cfg: omegaconf.DictConfig):
                 temperature (float): temperature for the softmax normalization.
                 queue_size (int): number of samples to hold in the queue.
                 epoch_queue_starts (int): epochs the queue starts.
-                freeze_prototypes_epochs (int): number of epochs during which the prototypes are frozen.
+                freeze_prototypes_epochs (int): number of epochs during which
+                    the prototypes are frozen.
         """
 
         super().__init__(cfg)