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

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

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

@@ -15,8 +15,8 @@ ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXp
 ring/algorithms/custom_joints/rsaddle_joint.py,sha256=QoMo6NXdYgA9JygSzBvr0eCdd3qKhUgCrGPNO2Qdxko,1200
 ring/algorithms/custom_joints/suntay.py,sha256=tOEGM304XciHO4pmvxr4faA4xXVO4N2HlPdFmXKbcrw,16726
 ring/algorithms/generator/__init__.py,sha256=bF-CW3x2x-o6KWESKy-DuxzZPh3UNSjJb_MaAcSHGsQ,277
-ring/algorithms/generator/base.py,sha256=rrhHg6lFPDJs72kXvzF15v1vzkaUTKtCnpcmWZONYA8,16847
-ring/algorithms/generator/batch.py,sha256=9yFxVv11hij-fJXGPxA3zEh1bE2_jrZk0R7kyGaiM5c,2551
+ring/algorithms/generator/base.py,sha256=jGQocoNZ5tkiMazBDCv-jD6FNYwebqn0_RgVFse49pg,16890
+ring/algorithms/generator/batch.py,sha256=P51UnAZl9TUF_eVq58VL1CsmPPStPHhRDdKjUyvu4EA,2652
 ring/algorithms/generator/finalize_fns.py,sha256=nY2RKiLbHriTkdec94lc4UGSZKd0v547MDNn4dr8I3E,10398
 ring/algorithms/generator/motion_artifacts.py,sha256=2VJbldVDbI3PSyboshIbtYvSAKzBBwGV7cQfYjqvluM,9167
 ring/algorithms/generator/pd_control.py,sha256=XJ_Gd5AkIRh-jBrMfQyMXjVwhx2gCNHznjzFbmAwhZs,5767
@@ -58,7 +58,7 @@ 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
 ring/ml/rnno_v1.py,sha256=2qE08OIvTJ5PvSxKpYGzGSrvEImWrdAT_qslZ7jP5tA,1372
-ring/ml/train.py,sha256=-6SzQKjIgktgRjaXKVg_1dqcBmAJggZSVwDnau1FnxI,10832
+ring/ml/train.py,sha256=Da89HxiqXC7xuX2ldpTrJStqKWN-6Vcpml4PPQuihN4,10989
 ring/ml/training_loop.py,sha256=CEokvPQuuk_WCd-J60ZDodJYcPVvyxLfgXDr_DnbzRI,3359
 ring/ml/params/0x13e3518065c21cd8.pickle,sha256=Zh2k1zK-TNxJl5F7nyTeQ9001qqRE_dfvaq1HWV287A,9355838
 ring/ml/params/0x1d76628065a71e0f.pickle,sha256=YTNVuvfw-nCRD9BH1PZYcR9uCFpNWDhw8Lc50eDn_EE,9351038
@@ -78,7 +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=wMKJ-eCJ4cHjisGODOZgDVG2r-XQjSANBQFfC05wpzo,2092
+ring/utils/dataloader_torch.py,sha256=bravdBqbkxxcQDieg6OnmArGwGcWpMI3MmNFwTCt0qg,3808
 ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
 ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=zRPfxYNesvgefkddd26oar6f9433LkMGkhp9dF3rPUs,1926
@@ -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.32.dist-info/METADATA,sha256=6bNRA4bhdmUOnqV8ZfR-tl5wbk-awzTE7qmLzDOI1xs,4251
-imt_ring-1.6.32.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-imt_ring-1.6.32.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.32.dist-info/RECORD,,
+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,,

ring/algorithms/generator/base.py

@@ -213,6 +213,8 @@ class RCMG:
                 )
                 save_fn(d, file)
                 i += 1
+            # cleanup
+            del data
 
         gens, n_calls = self._generators_ncalls(sizes)
         batch.generators_eager(gens, n_calls, callback, seed, self._disable_tqdm)

ring/algorithms/generator/batch.py

@@ -1,3 +1,4 @@
+import gc
 from typing import Callable
 
 import jax
@@ -83,4 +84,8 @@ def generators_eager(
 
             sample_flat, _ = jax.tree_util.tree_flatten(sample)
             size = 1 if len(sample_flat) == 0 else sample_flat[0].shape[0]
-            callback([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
+            callback([jax.tree_map(lambda a: a[i].copy(), sample) for i in range(size)])
+
+            # cleanup
+            del sample, sample_flat
+            gc.collect()

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

ring/utils/dataloader_torch.py

@@ -1,8 +1,9 @@
 import os
-from typing import Optional
+from typing import Any, Optional
 import warnings
 
 import jax
+import numpy as np
 import torch
 from torch.utils.data import DataLoader
 from torch.utils.data import Dataset
@@ -12,7 +13,7 @@ from ring.utils import parse_path
 from ring.utils import pickle_load
 
 
-class FolderOfPickleFilesDataset(Dataset):
+class FolderOfFilesDataset(Dataset):
     def __init__(self, path, transform=None):
         self.files = self.listdir(path)
         self.transform = transform
@@ -22,7 +23,7 @@ class FolderOfPickleFilesDataset(Dataset):
         return self.N
 
     def __getitem__(self, idx: int):
-        element = pickle_load(self.files[idx])
+        element = self._load_file(self.files[idx])
         if self.transform is not None:
             element = self.transform(element)
         return element
@@ -31,6 +32,10 @@ class FolderOfPickleFilesDataset(Dataset):
     def listdir(path: str) -> list:
         return [parse_path(path, file) for file in os.listdir(path)]
 
+    @staticmethod
+    def _load_file(file_path: str) -> Any:
+        return pickle_load(file_path)
+
 
 def dataset_to_generator(
     dataset: Dataset,
@@ -84,3 +89,60 @@ def _get_number_of_logical_cores() -> int:
         )
         N = 0
     return N
+
+
+class MultiDataset(Dataset):
+    def __init__(self, datasets, transform=None):
+        """
+        Args:
+            datasets: A list of datasets to sample from.
+            transform: A function that takes N items (one from each dataset) and combines them.
+        """  # noqa: E501
+        self.datasets = datasets
+        self.transform = transform
+
+    def __len__(self):
+        # Length is defined by the smallest dataset in the list
+        return min(len(ds) for ds in self.datasets)
+
+    def __getitem__(self, idx):
+        sampled_items = [ds[idx] for ds in self.datasets]
+
+        if self.transform:
+            # Apply the transformation to all sampled items
+            return self.transform(*sampled_items)
+
+        return tuple(sampled_items)
+
+
+class ShuffledDataset(Dataset):
+    def __init__(self, dataset):
+        """
+        Wrapper that shuffles the dataset indices once.
+
+        Args:
+            dataset (Dataset): The original dataset to shuffle.
+        """
+        self.dataset = dataset
+        self.shuffled_indices = np.random.permutation(
+            len(dataset)
+        )  # Shuffle indices once
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        """
+        Returns the data at the shuffled index.
+
+        Args:
+            idx (int): Index in the shuffled dataset.
+        """
+        original_idx = self.shuffled_indices[idx]
+        return self.dataset[original_idx]
+
+
+def dataset_to_Xy(ds: Dataset):
+    return dataset_to_generator(ds, batch_size=len(ds), shuffle=False, num_workers=0)(
+        None
+    )