imt-ring 1.6.30__tar.gz → 1.6.31__tar.gz

{imt_ring-1.6.30 → imt_ring-1.6.31}/PKG-INFO

@@ -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 → imt_ring-1.6.31}/pyproject.toml

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

{imt_ring-1.6.30 → imt_ring-1.6.31}/src/imt_ring.egg-info/PKG-INFO

@@ -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 → imt_ring-1.6.31}/src/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