imt-ring 1.6.33__tar.gz → 1.6.35__tar.gz

{imt_ring-1.6.33 → imt_ring-1.6.35}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.33
+Version: 1.6.35
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.33 → imt_ring-1.6.35}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imt-ring"
-version = "1.6.33"
+version = "1.6.35"
 authors = [
   { name="Simon Bachhuber", email="simon.bachhuber@fau.de" },
 ]

{imt_ring-1.6.33 → imt_ring-1.6.35}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.33
+Version: 1.6.35
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.33 → imt_ring-1.6.35}/src/ring/ml/callbacks.py

@@ -388,6 +388,14 @@ class TimingKillRunCallback(training_loop.TrainingLoopCallback):
 
 
 class CheckpointCallback(training_loop.TrainingLoopCallback):
+    def __init__(
+        self,
+        checkpoint_every: Optional[int] = None,
+        checkpoint_folder: str = "~/.ring_checkpoints",
+    ):
+        self.checkpoint_every = checkpoint_every
+        self.checkpoint_folder = checkpoint_folder
+
     def after_training_step(
         self,
         i_episode: int,
@@ -401,18 +409,26 @@ class CheckpointCallback(training_loop.TrainingLoopCallback):
         self.params = params
         self.opt_state = opt_state
 
+        if self.checkpoint_every is not None and (
+            (i_episode % self.checkpoint_every) == 0
+        ):
+            self._create_checkpoint()
+
+    def _create_checkpoint(self):
+        path = parse_path(
+            self.checkpoint_folder, ml_utils.unique_id(), extension="pickle"
+        )
+        data = {"params": self.params, "opt_state": self.opt_state}
+        pickle_save(
+            obj=jax.device_get(data),
+            path=path,
+            overwrite=True,
+        )
+
     def close(self):
         # only checkpoint if run has been killed
         if training_loop.recv_kill_run_signal():
-            path = parse_path(
-                "~/.ring_checkpoints", ml_utils.unique_id(), extension="pickle"
-            )
-            data = {"params": self.params, "opt_state": self.opt_state}
-            pickle_save(
-                obj=jax.device_get(data),
-                path=path,
-                overwrite=True,
-            )
+            self._create_checkpoint()
 
 
 class WandbKillRun(training_loop.TrainingLoopCallback):

{imt_ring-1.6.33 → imt_ring-1.6.35}/src/ring/ml/train.py

@@ -39,6 +39,7 @@ def _build_step_fn(
     filter: ml_base.AbstractFilter,
     optimizer,
     tbp,
+    skip_first_tbp_batch,
 ):
     """Build step function that optimizes filter parameters based on `metric_fn`.
     `initial_state` has shape (pmap, vmap, state_dim)"""
@@ -89,6 +90,8 @@ def _build_step_fn(
         ):
             (loss, state), grads = pmapped_loss_fn(params, state, X_tbp, y_tbp)
             debug_grads.append(grads)
+            if skip_first_tbp_batch and i == 0:
+                continue
             state = jax.lax.stop_gradient(state)
             params, opt_state = apply_grads(grads, params, opt_state)
 
@@ -119,6 +122,7 @@ def train_fn(
     loss_fn: LOSS_FN = _default_loss_fn,
     metrices: Optional[METRICES] = _default_metrices,
     link_names: Optional[list[str]] = None,
+    skip_first_tbp_batch: bool = False,
 ) -> bool:
     """Trains RNNO
 
@@ -161,10 +165,7 @@ def train_fn(
         opt_state = optimizer.init(filter_params)
 
     step_fn = _build_step_fn(
-        loss_fn,
-        filter,
-        optimizer,
-        tbp=tbp,
+        loss_fn, filter, optimizer, tbp=tbp, skip_first_tbp_batch=skip_first_tbp_batch
     )
 
     # always log, because we also want `i_epsiode` to be logged in wandb