feat: add pypots.random.set_random_seed();

pypots/tests/test_utils.py

@@ -9,7 +9,10 @@
 import shutil
 import unittest
 
+import torch
+
 from pypots.utils.logging import Logger
+from pypots.utils.random import set_random_seed
 
 
 class TestLogging(unittest.TestCase):
@@ -46,5 +49,25 @@ def test_saving_log_into_file(self):
  shutil.rmtree("test_log", ignore_errors=True)
 
 
+class TestRandom(unittest.TestCase):
+ def test_set_random_seed(self):
+ random_state1 = torch.get_rng_state()
+ torch.rand(
+ 1, 3
+ ) # randomly generate something, the random state will be reset, so two states should be varying
+ random_state2 = torch.get_rng_state()
+ assert not torch.equal(
+ random_state1, random_state2
+ ), "The random seed hasn't set, so two random states should be different."
+
+ set_random_seed(26)
+ random_state1 = torch.get_rng_state()
+ set_random_seed(26)
+ random_state2 = torch.get_rng_state()
+ assert torch.equal(
+ random_state1, random_state2
+ ), "The random seed has been set, two random states are not the same."
+
+
 if __name__ == "__main__":
  unittest.main()

pypots/utils/random.py

@@ -0,0 +1,29 @@
+"""
+Transformer model for time-series imputation.
+"""
+
+# Created by Wenjie Du <[email protected]>
+# License: GLP-v3
+
+import numpy as np
+import torch
+from pypots.utils.logging import logger
+
+RANDOM_SEED = 2204
+
+
+def set_random_seed(random_seed: int = RANDOM_SEED):
+ """Manually set the random state to make PyPOTS output reproducible results.
+
+ Parameters
+ ----------
+ random_seed : int, default = RANDOM_SEED,
+ The seed to be set for generating random numbers in PyPOTS.
+
+ """
+
+ np.random.seed(RANDOM_SEED)
+ torch.manual_seed(random_seed)
+ logger.info(
+ f"Done. Have already set the random seed as {random_seed} for numpy and pytorch."
+ )