imt-ring 1.6.30__py3-none-any.whl → 1.6.31__py3-none-any.whl

{imt_ring-1.6.30.dist-info → imt_ring-1.6.31.dist-info}/METADATA

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

@@ -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=DR2uUiA9x49_6EBjnbVLfWu7GBX7wtKjgHSIlF80HO0,1502
+ring/utils/dataloader_torch.py,sha256=wMKJ-eCJ4cHjisGODOZgDVG2r-XQjSANBQFfC05wpzo,2092
 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.30.dist-info/METADATA,sha256=CC1Tqt_uwTMJZiSWoxBCLjkc3J1hWsNhPTJNsxamM0U,4251
-imt_ring-1.6.30.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-imt_ring-1.6.30.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.30.dist-info/RECORD,,
+imt_ring-1.6.31.dist-info/METADATA,sha256=ISL0fShgxIGskumWa3mtqCBcOvOfEtJc4XE8Rt-EJCA,4251
+imt_ring-1.6.31.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+imt_ring-1.6.31.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.31.dist-info/RECORD,,

ring/utils/dataloader_torch.py

@@ -1,4 +1,6 @@
 import os
+from typing import Optional
+import warnings
 
 import jax
 import torch
@@ -35,15 +37,20 @@ def dataset_to_generator(
     batch_size: int,
     shuffle=True,
     seed: int = 1,
+    num_workers: Optional[int] = None,
     **kwargs,
 ):
     torch.manual_seed(seed)
 
+    if num_workers is None:
+        num_workers = _get_number_of_logical_cores()
+
     dl = DataLoader(
         dataset,
         batch_size=batch_size,
         shuffle=shuffle,
-        multiprocessing_context="spawn" if kwargs.get("num_workers", 0) > 0 else None,
+        multiprocessing_context="spawn" if num_workers > 0 else None,
+        num_workers=num_workers,
         **kwargs,
     )
     dl_iter = iter(dl)
@@ -60,3 +67,20 @@ def dataset_to_generator(
             return to_numpy(next(dl_iter))
 
     return generator
+
+
+def _get_number_of_logical_cores() -> int:
+    N = None
+    if hasattr(os, "sched_getaffinity"):
+        try:
+            N = len(os.sched_getaffinity(0))
+        except Exception:
+            pass
+    if N is None:
+        N = os.cpu_count()
+    if N is None:
+        warnings.warn(
+            "Could not automatically set the `num_workers` variable, defaults to `0`"
+        )
+        N = 0
+    return N