Merge pull request #36 from BirkhoffG/reproducibility

BirkhoffG · web-flow · commit fcec10fb024f · 2024-10-18T15:34:06.000-04:00
Implement manual_seed to set the global seed work for Reproducibility
diff --git a/jax_dataloader/_modidx.py b/jax_dataloader/_modidx.py
@@ -113,7 +113,9 @@
                                                                                           'jax_dataloader/tests.py'),
                                       'jax_dataloader.tests.test_shuffle': ('tests.html#test_shuffle', 'jax_dataloader/tests.py'),
                                       'jax_dataloader.tests.test_shuffle_drop_last': ( 'tests.html#test_shuffle_drop_last',
-                                                                                       'jax_dataloader/tests.py')},
+                                                                                       'jax_dataloader/tests.py'),
+                                      'jax_dataloader.tests.test_shuffle_reproducible': ( 'tests.html#test_shuffle_reproducible',
+                                                                                          'jax_dataloader/tests.py')},
             'jax_dataloader.utils': { 'jax_dataloader.utils.Config': ('utils.html#config', 'jax_dataloader/utils.py'),
                                       'jax_dataloader.utils.Config.default': ('utils.html#config.default', 'jax_dataloader/utils.py'),
                                       'jax_dataloader.utils.asnumpy': ('utils.html#asnumpy', 'jax_dataloader/utils.py'),
@@ -125,4 +127,5 @@
                                                                                    'jax_dataloader/utils.py'),
                                       'jax_dataloader.utils.get_config': ('utils.html#get_config', 'jax_dataloader/utils.py'),
                                       'jax_dataloader.utils.has_pytorch_tensor': ( 'utils.html#has_pytorch_tensor',
-                                                                                   'jax_dataloader/utils.py')}}}
+                                                                                   'jax_dataloader/utils.py'),
+                                      'jax_dataloader.utils.manual_seed': ('utils.html#manual_seed', 'jax_dataloader/utils.py')}}}
diff --git a/jax_dataloader/loaders/torch.py b/jax_dataloader/loaders/torch.py
@@ -5,7 +5,7 @@
 from ..imports import *
 from . import BaseDataLoader
 from ..datasets import Dataset, ArrayDataset, JAXDataset
-from ..utils import check_pytorch_installed
+from ..utils import check_pytorch_installed, get_config
 from ..tests import *
 from jax.tree_util import tree_map
 import warnings
@@ -53,13 +53,20 @@ def __init__(
         super().__init__(dataset, batch_size, shuffle, drop_last)
         check_pytorch_installed()
         from torch.utils.data import BatchSampler, RandomSampler, SequentialSampler
+        import torch
 
         if 'sampler' in kwargs:
             warnings.warn("`sampler` is currently not supported. We will ignore it and use `shuffle` instead.")
             del kwargs['sampler']
 
+        # convert to torch dataset
         dataset = to_torch_dataset(dataset)
-        sampler = RandomSampler(dataset) if shuffle else SequentialSampler(dataset)
+        # init batch sampler
+        generator = torch.Generator().manual_seed(get_config().global_seed)
+        if shuffle: 
+            sampler = RandomSampler(dataset, generator=generator)
+        else:       
+            sampler = SequentialSampler(dataset)
         batch_sampler = BatchSampler(sampler, batch_size=batch_size, drop_last=drop_last)
 
         self.dataloader = torch_data.DataLoader(
diff --git a/jax_dataloader/tests.py b/jax_dataloader/tests.py
@@ -4,9 +4,10 @@
 from __future__ import print_function, division, annotations
 from .imports import *
 from .datasets import ArrayDataset
+import jax_dataloader as jdl
 
 # %% auto 0
-__all__ = ['test_dataloader']
+__all__ = ['test_shuffle_reproducible', 'test_dataloader']
 
 # %% ../nbs/tests.ipynb 3
 def get_batch(batch):
@@ -82,6 +83,31 @@ def test_shuffle_drop_last(cls, ds, batch_size: int, feats, labels):
         assert len(_X) == len(X_list) * batch_size
 
 # %% ../nbs/tests.ipynb 8
+def test_shuffle_reproducible(cls, ds, batch_size: int, feats, labels):
+    """Test that the shuffle is reproducible"""
+    def _iter_dataloader(dataloader):
+        X_list, Y_list = [], []
+        for batch in dataloader:
+            x, y = get_batch(batch)
+            X_list.append(x)
+            Y_list.append(y)
+        return X_list, Y_list
+
+    # Test that the shuffle is reproducible
+    jdl.manual_seed(0)
+    dl_1 = cls(ds, batch_size=batch_size, shuffle=True, drop_last=False)
+    X_list_1, Y_list_1 = _iter_dataloader(dl_1)
+    dl_2 = cls(ds, batch_size=batch_size, shuffle=True, drop_last=False)
+    X_list_2, Y_list_2 = _iter_dataloader(dl_2)
+    assert jnp.array_equal(jnp.concatenate(X_list_1), jnp.concatenate(X_list_2))
+
+    # Test that the shuffle is different if the seed is different
+    jdl.manual_seed(1234)
+    dl_3 = cls(ds, batch_size=batch_size, shuffle=True, drop_last=False)
+    X_list_3, Y_list_3 = _iter_dataloader(dl_3)
+    assert not jnp.array_equal(jnp.concatenate(X_list_1), jnp.concatenate(X_list_3))
+
+# %% ../nbs/tests.ipynb 9
 def test_dataloader(cls, ds_type='jax', samples=1000, batch_size=12):
     feats = np.arange(samples).repeat(10).reshape(samples, 10)
     labels = np.arange(samples).reshape(samples, 1)
@@ -102,3 +128,4 @@ def test_dataloader(cls, ds_type='jax', samples=1000, batch_size=12):
     test_no_shuffle_drop_last(cls, ds, batch_size, feats, labels)
     test_shuffle(cls, ds, batch_size, feats, labels)
     test_shuffle_drop_last(cls, ds, batch_size, feats, labels)
+    test_shuffle_reproducible(cls, ds, batch_size, feats, labels)
diff --git a/jax_dataloader/utils.py b/jax_dataloader/utils.py
@@ -7,8 +7,8 @@
 import collections
 
 # %% auto 0
-__all__ = ['Config', 'get_config', 'check_pytorch_installed', 'has_pytorch_tensor', 'check_hf_installed', 'check_tf_installed',
-           'asnumpy']
+__all__ = ['Config', 'get_config', 'manual_seed', 'check_pytorch_installed', 'has_pytorch_tensor', 'check_hf_installed',
+           'check_tf_installed', 'asnumpy']
 
 # %% ../nbs/utils.ipynb 6
 @dataclass
@@ -28,15 +28,20 @@ def default(cls) -> Config:
 def get_config() -> Config:
     return main_config
 
-# %% ../nbs/utils.ipynb 10
+# %% ../nbs/utils.ipynb 9
+def manual_seed(seed: int):
+    """Set the seed for the library"""
+    main_config.global_seed = seed
+
+# %% ../nbs/utils.ipynb 12
 def check_pytorch_installed():
     if torch_data is None:
         raise ModuleNotFoundError("`pytorch` library needs to be installed. "
             "Try `pip install torch`. Please refer to pytorch documentation for details: "
             "https://pytorch.org/get-started/.")
 
 
-# %% ../nbs/utils.ipynb 12
+# %% ../nbs/utils.ipynb 14
 def has_pytorch_tensor(batch) -> bool:
     if isinstance(batch[0], torch.Tensor):
         return True
@@ -46,21 +51,21 @@ def has_pytorch_tensor(batch) -> bool:
     else:
         return False
 
-# %% ../nbs/utils.ipynb 13
+# %% ../nbs/utils.ipynb 15
 def check_hf_installed():
     if hf_datasets is None:
         raise ModuleNotFoundError("`datasets` library needs to be installed. "
             "Try `pip install datasets`. Please refer to huggingface documentation for details: "
             "https://huggingface.co/docs/datasets/installation.html.")
 
-# %% ../nbs/utils.ipynb 15
+# %% ../nbs/utils.ipynb 17
 def check_tf_installed():
     if tf is None:
         raise ModuleNotFoundError("`tensorflow` library needs to be installed. "
             "Try `pip install tensorflow`. Please refer to tensorflow documentation for details: "
             "https://www.tensorflow.org/install/pip.")
 
-# %% ../nbs/utils.ipynb 18
+# %% ../nbs/utils.ipynb 20
 def asnumpy(x) -> np.ndarray:
     if isinstance(x, np.ndarray):
         return x
diff --git a/nbs/loader.torch.ipynb b/nbs/loader.torch.ipynb
@@ -27,25 +27,14 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2023-12-26 09:48:50.763338: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
-      "2023-12-26 09:48:50.763415: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
-      "2023-12-26 09:48:50.777782: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
-      "2023-12-26 09:48:52.348997: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "#| export\n",
     "from __future__ import print_function, division, annotations\n",
     "from jax_dataloader.imports import *\n",
     "from jax_dataloader.loaders import BaseDataLoader\n",
     "from jax_dataloader.datasets import Dataset, ArrayDataset, JAXDataset\n",
-    "from jax_dataloader.utils import check_pytorch_installed\n",
+    "from jax_dataloader.utils import check_pytorch_installed, get_config\n",
     "from jax_dataloader.tests import *\n",
     "from jax.tree_util import tree_map\n",
     "import warnings\n"
@@ -129,13 +118,20 @@
     "        super().__init__(dataset, batch_size, shuffle, drop_last)\n",
     "        check_pytorch_installed()\n",
     "        from torch.utils.data import BatchSampler, RandomSampler, SequentialSampler\n",
+    "        import torch\n",
     "\n",
     "        if 'sampler' in kwargs:\n",
     "            warnings.warn(\"`sampler` is currently not supported. We will ignore it and use `shuffle` instead.\")\n",
     "            del kwargs['sampler']\n",
     "\n",
+    "        # convert to torch dataset\n",
     "        dataset = to_torch_dataset(dataset)\n",
-    "        sampler = RandomSampler(dataset) if shuffle else SequentialSampler(dataset)\n",
+    "        # init batch sampler\n",
+    "        generator = torch.Generator().manual_seed(get_config().global_seed)\n",
+    "        if shuffle: \n",
+    "            sampler = RandomSampler(dataset, generator=generator)\n",
+    "        else:       \n",
+    "            sampler = SequentialSampler(dataset)\n",
     "        batch_sampler = BatchSampler(sampler, batch_size=batch_size, drop_last=drop_last)\n",
     "\n",
     "        self.dataloader = torch_data.DataLoader(\n",
diff --git a/nbs/tests.ipynb b/nbs/tests.ipynb
@@ -32,7 +32,8 @@
     "#| export\n",
     "from __future__ import print_function, division, annotations\n",
     "from jax_dataloader.imports import *\n",
-    "from jax_dataloader.datasets import ArrayDataset"
+    "from jax_dataloader.datasets import ArrayDataset\n",
+    "import jax_dataloader as jdl"
    ]
   },
   {
@@ -143,6 +144,38 @@
     "        assert len(_X) == len(X_list) * batch_size"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "def test_shuffle_reproducible(cls, ds, batch_size: int, feats, labels):\n",
+    "    \"\"\"Test that the shuffle is reproducible\"\"\"\n",
+    "    def _iter_dataloader(dataloader):\n",
+    "        X_list, Y_list = [], []\n",
+    "        for batch in dataloader:\n",
+    "            x, y = get_batch(batch)\n",
+    "            X_list.append(x)\n",
+    "            Y_list.append(y)\n",
+    "        return X_list, Y_list\n",
+    "\n",
+    "    # Test that the shuffle is reproducible\n",
+    "    jdl.manual_seed(0)\n",
+    "    dl_1 = cls(ds, batch_size=batch_size, shuffle=True, drop_last=False)\n",
+    "    X_list_1, Y_list_1 = _iter_dataloader(dl_1)\n",
+    "    dl_2 = cls(ds, batch_size=batch_size, shuffle=True, drop_last=False)\n",
+    "    X_list_2, Y_list_2 = _iter_dataloader(dl_2)\n",
+    "    assert jnp.array_equal(jnp.concatenate(X_list_1), jnp.concatenate(X_list_2))\n",
+    "\n",
+    "    # Test that the shuffle is different if the seed is different\n",
+    "    jdl.manual_seed(1234)\n",
+    "    dl_3 = cls(ds, batch_size=batch_size, shuffle=True, drop_last=False)\n",
+    "    X_list_3, Y_list_3 = _iter_dataloader(dl_3)\n",
+    "    assert not jnp.array_equal(jnp.concatenate(X_list_1), jnp.concatenate(X_list_3))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -169,7 +202,8 @@
     "    test_no_shuffle(cls, ds, batch_size, feats, labels)\n",
     "    test_no_shuffle_drop_last(cls, ds, batch_size, feats, labels)\n",
     "    test_shuffle(cls, ds, batch_size, feats, labels)\n",
-    "    test_shuffle_drop_last(cls, ds, batch_size, feats, labels)"
+    "    test_shuffle_drop_last(cls, ds, batch_size, feats, labels)\n",
+    "    test_shuffle_reproducible(cls, ds, batch_size, feats, labels)"
    ]
   },
   {
diff --git a/nbs/utils.ipynb b/nbs/utils.ipynb
@@ -99,6 +99,28 @@
     "    return main_config"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "def manual_seed(seed: int):\n",
+    "    \"\"\"Set the seed for the library\"\"\"\n",
+    "    main_config.global_seed = seed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "manual_seed(11)\n",
+    "assert get_config().global_seed == 11"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
