revert unneeded changes in solver

Author: stes

cebra/solver/base.py

@@ -31,9 +31,7 @@
 """
 
 import abc
-import logging
 import os
-import time
 from typing import Callable, Dict, List, Literal, Optional
 
 import literate_dataclasses as dataclasses
@@ -72,10 +70,6 @@ class Solver(abc.ABC, cebra.io.HasDevice):
     optimizer: torch.optim.Optimizer
     history: List = dataclasses.field(default_factory=list)
     decode_history: List = dataclasses.field(default_factory=list)
-    metadata: Dict = dataclasses.field(default_factory=lambda: ({
-        "timestamp": None,
-        "batches_seen": None,
-    }))
     log: Dict = dataclasses.field(default_factory=lambda: ({
         "pos": [],
         "neg": [],
@@ -84,8 +78,6 @@ class Solver(abc.ABC, cebra.io.HasDevice):
     }))
     tqdm_on: bool = True
 
-    #metrics: MetricCollection = None
-
     def __post_init__(self):
         cebra.io.HasDevice.__init__(self)
         self.best_loss = float("inf")
@@ -105,7 +97,6 @@ def state_dict(self) -> dict:
             "loss": torch.tensor(self.history),
             "decode": self.decode_history,
             "criterion": self.criterion.state_dict(),
-            "metadata": self.metadata,
             "version": cebra.__version__,
             "log": self.log,
         }
@@ -120,7 +111,7 @@ def load_state_dict(self, state_dict: dict, strict: bool = True):
                 to partially load the state for all given keys.
         """
 
-        def _contains(key, strict=strict):
+        def _contains(key):
             if key in state_dict:
                 return True
             elif strict:
@@ -146,9 +137,6 @@ def _get(key):
             self.decode_history = _get("decode")
         if _contains("log"):
             self.log = _get("log")
-        # NOTE(stes): Added in CEBRA 0.6.0
-        if _contains("metadata", strict=False):
-            self.metadata = _get("metadata")
 
     @property
     def num_parameters(self) -> int:
@@ -163,26 +151,21 @@ def parameters(self):
         for parameter in self.criterion.parameters():
             yield parameter
 
-    def _get_loader(self, loader, **kwargs):
-        return ProgressBar(loader=loader,
-                           log_format="tqdm" if self.tqdm_on else "off",
-                           **kwargs)
-
-    def _update_metadata(self, num_steps):
-        self.metadata["timestamp"] = time.time()
-        self.metadata["batches_seen"] = num_steps
+    def _get_loader(self, loader):
+        return ProgressBar(
+            loader,
+            "tqdm" if self.tqdm_on else "off",
+        )
 
     def fit(self,
             loader: cebra.data.Loader,
             valid_loader: cebra.data.Loader = None,
             *,
             save_frequency: int = None,
             valid_frequency: int = None,
-            log_frequency: int = None,
             decode: bool = False,
             logdir: str = None,
-            save_hook: Callable[[int, "Solver"], None] = None,
-            logger: logging.Logger = None):
+            save_hook: Callable[[int, "Solver"], None] = None):
         """Train model for the specified number of steps.
 
         Args:
@@ -192,27 +175,20 @@ def fit(self,
             save_frequency: If not `None`, the frequency for automatically saving model checkpoints
                 to `logdir`.
             valid_frequency: The frequency for running validation on the ``valid_loader`` instance.
-            log_frequency: TODO
             logdir:  The logging directory for writing model checkpoints. The checkpoints
                 can be read again using the `solver.load` function, or manually via loading the
                 state dict.
-            logger: TODO
 
         TODO:
             * Refine the API here. Drop the validation entirely, and implement this via a hook?
         """
 
         self.to(loader.device)
 
-        iterator = self._get_loader(loader,
-                                    logger=logger,
-                                    log_frequency=log_frequency)
-
         iterator = self._get_loader(loader)
         self.model.train()
         for num_steps, batch in iterator:
             stats = self.step(batch)
-            self._update_metadata(num_steps)
             iterator.set_description(stats)
 
             if save_frequency is None:
@@ -476,15 +452,11 @@ def step(self, batch: cebra.data.Batch) -> dict:
         self.optimizer.step()
         self.history.append(loss.item())
 
-        stats = dict(
+        return dict(
             behavior_pos=behavior_align.item(),
             behavior_neg=behavior_uniform.item(),
             behavior_total=behavior_loss.item(),
             time_pos=time_align.item(),
             time_neg=time_uniform.item(),
             time_total=time_loss.item(),
         )
-
-        for key, value in stats.items():
-            self.log[key].append(value)
-        return stats