imt-ring 1.6.34__py3-none-any.whl → 1.6.35__py3-none-any.whl

{imt_ring-1.6.34.dist-info → imt_ring-1.6.35.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.34
+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.34.dist-info → imt_ring-1.6.35.dist-info}/RECORD

@@ -53,7 +53,7 @@ ring/io/xml/test_to_xml.py,sha256=NGn4VSiFdwhYN5YTBduWMiY9B5dwtxZhCQAR_PXeqKU,94
 ring/io/xml/to_xml.py,sha256=Wo4iySLw9nM-iVW42AGvMRqjtU2qRc2FD_Zlc7w1IrE,3438
 ring/ml/__init__.py,sha256=nbh48gaswWeY4S4vT1sply_3ROj2DQ7agjoLR4Ho3T8,1517
 ring/ml/base.py,sha256=lfwEZLBDglOSRWChUHoH1kezefhttPV9TMEpNIqsMNw,9972
-ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
+ring/ml/callbacks.py,sha256=oCPXl4_Zcw3g0KRgyyUDmdiGxV0phnDVc_t8rEG4Lls,13737
 ring/ml/ml_utils.py,sha256=M--qkXRnhU7tHvgfTHfT9gyY0nhj3zMGEaK0X0drFLs,10915
 ring/ml/optimizer.py,sha256=TZF0_LmnewzmGVso-zIQJtpWguUW0fW3HeRpIdG_qoI,4763
 ring/ml/ringnet.py,sha256=mef7jyN2QcApJmQGH3HYZyTV-00q8YpsYOKhW0-ku1k,8973
@@ -86,7 +86,7 @@ ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3
 ring/utils/utils.py,sha256=tJaWXLGOTwkxJQj2l23dX97wO3aZYhM2qd7eNuMRs84,6907
 ring/utils/register_gym_envs/__init__.py,sha256=PtPIRBQJ16339xZ9G9VpvqrvcGbQ_Pk_SUz4tQPa9nQ,94
 ring/utils/register_gym_envs/saddle.py,sha256=tA5CyW_akSXyDm0xJ83CtOrUMVElH0f9vZtEDDJQalI,4422
-imt_ring-1.6.34.dist-info/METADATA,sha256=D7FXQFI8b4iXaJiWkM4MvHNqzncvFh_wn4UpVK8iqMs,4251
-imt_ring-1.6.34.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-imt_ring-1.6.34.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.34.dist-info/RECORD,,
+imt_ring-1.6.35.dist-info/METADATA,sha256=mwRs6Hzb4M39Ld4PRamHOYPxtbu-Ug2kOLE03M1l8hI,4251
+imt_ring-1.6.35.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+imt_ring-1.6.35.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.35.dist-info/RECORD,,

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):