imt-ring 1.6.17__py3-none-any.whl → 1.6.19__py3-none-any.whl

{imt_ring-1.6.17.dist-info → imt_ring-1.6.19.dist-info}/METADATA

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

@@ -54,7 +54,7 @@ ring/io/xml/to_xml.py,sha256=fohb-jWMf2cxVdT5dmknsGyrNMseICSbKEz_urbaWbQ,3407
 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/ml_utils.py,sha256=1GXJfeoXbwCbRdYA2np3CbJpSupaw4eyf3quh9y4BO0,6462
+ring/ml/ml_utils.py,sha256=xqy9BnLy8IKVqkFS9mlZsGJXSbThI9zZxZ5rhl8LSI8,7144
 ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
 ring/ml/ringnet.py,sha256=mef7jyN2QcApJmQGH3HYZyTV-00q8YpsYOKhW0-ku1k,8973
 ring/ml/rnno_v1.py,sha256=2qE08OIvTJ5PvSxKpYGzGSrvEImWrdAT_qslZ7jP5tA,1372
@@ -78,6 +78,7 @@ ring/utils/backend.py,sha256=cKSi9sB59texqKzNVASTDczGKLCBL8VVDiP7TNdj41k,1294
 ring/utils/batchsize.py,sha256=FbOii7MDP4oPZd9GJOKehFatfnb6WZ0b9z349iZYs1A,1786
 ring/utils/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
 ring/utils/dataloader.py,sha256=2CcsbUY2AZs8LraS5HTJXlEseuF-1gKmfyBkSsib-tE,3748
+ring/utils/dataloader_torch.py,sha256=DR2uUiA9x49_6EBjnbVLfWu7GBX7wtKjgHSIlF80HO0,1502
 ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
 ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=zRPfxYNesvgefkddd26oar6f9433LkMGkhp9dF3rPUs,1926
@@ -85,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.17.dist-info/METADATA,sha256=j0IKIyc6qAgz9059Z4-b46hi7qTn_ffI7oLw_OrD_Tk,3833
-imt_ring-1.6.17.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-imt_ring-1.6.17.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.17.dist-info/RECORD,,
+imt_ring-1.6.19.dist-info/METADATA,sha256=BZiSbypG96pTKHUSc1NDZYcyiRf1-ynsnu6rAsVdiRU,3833
+imt_ring-1.6.19.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+imt_ring-1.6.19.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.19.dist-info/RECORD,,

ring/ml/ml_utils.py

@@ -3,6 +3,7 @@ from functools import partial
 import os
 from pathlib import Path
 import pickle
+import shutil
 import time
 from typing import Optional, Protocol
 import warnings
@@ -107,7 +108,7 @@ class WandbLogger(MixinLogger):
         wandb.log(data)
 
     def log_params(self, path: str):
-        wandb.save(path, policy="now")
+        self.wandb_save(path)
 
     def log_video(
         self,
@@ -117,7 +118,7 @@ class WandbLogger(MixinLogger):
         step: Optional[int] = None,
     ):
         # TODO >>>
-        wandb.save(path, policy="now")
+        self.wandb_save(path)
         return
         # <<<
         data = {"video": wandb.Video(path, caption=caption, fps=fps)}
@@ -127,10 +128,10 @@ class WandbLogger(MixinLogger):
 
     def log_image(self, path: str, caption: Optional[str] = None):
         # wandb.log({"image": wandb.Image(path, caption=caption)})
-        wandb.save(path, policy="now")
+        self.wandb_save(path)
 
     def log_txt(self, path: str, wait: bool = True):
-        wandb.save(path, policy="now")
+        self.wandb_save(path)
         # TODO: `wandb` is not async at all?
         if wait:
             time.sleep(3)
@@ -138,6 +139,23 @@ class WandbLogger(MixinLogger):
     def close(self):
         wandb.run.finish()
 
+    @staticmethod
+    def wandb_save(path):
+        if wandb.run is not None and wandb.run.settings._offline:
+            # Create a dedicated directory in the WandB run directory to store copies
+            # of files
+            destination_dir = os.path.join(wandb.run.dir, "copied_files")
+            os.makedirs(destination_dir, exist_ok=True)
+
+            # Copy the file to this new location
+            copied_file_path = os.path.join(destination_dir, os.path.basename(path))
+            shutil.copy2(path, copied_file_path)
+
+            # Use wandb.save to save the copied file (now a true copy)
+            wandb.save(copied_file_path)
+        else:
+            wandb.save(path, policy="now")
+
 
 def _flatten_convert_filter_nested_dict(
     metrices: NestedDict, filter_nan_inf: bool = True

ring/utils/dataloader_torch.py

@@ -0,0 +1,62 @@
+import os
+
+import jax
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from tree_utils import PyTree
+
+from ring.utils import parse_path
+from ring.utils import pickle_load
+
+
+class FolderOfPickleFilesDataset(Dataset):
+    def __init__(self, path, transform=None):
+        self.files = self.listdir(path)
+        self.transform = transform
+        self.N = len(self.files)
+
+    def __len__(self):
+        return self.N
+
+    def __getitem__(self, idx: int):
+        element = pickle_load(self.files[idx])
+        if self.transform is not None:
+            element = self.transform(element)
+        return element
+
+    @staticmethod
+    def listdir(path: str) -> list:
+        return [parse_path(path, file) for file in os.listdir(path)]
+
+
+def dataset_to_generator(
+    dataset: Dataset,
+    batch_size: int,
+    shuffle=True,
+    seed: int = 1,
+    **kwargs,
+):
+    torch.manual_seed(seed)
+
+    dl = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        multiprocessing_context="spawn" if kwargs.get("num_workers", 0) > 0 else None,
+        **kwargs,
+    )
+    dl_iter = iter(dl)
+
+    def to_numpy(tree: PyTree[torch.Tensor]):
+        return jax.tree_map(lambda tensor: tensor.numpy(), tree)
+
+    def generator(_):
+        nonlocal dl, dl_iter
+        try:
+            return to_numpy(next(dl_iter))
+        except StopIteration:
+            dl_iter = iter(dl)
+            return to_numpy(next(dl_iter))
+
+    return generator