fix(reproducibility): add opt-in strict determinism across trainers

README.md

@@ -77,6 +77,12 @@ python engiopt/cgan_cnn_2d/cgan_cnn_2d.py --problem-id "beams2d" --track --wandb
 
 This will run a CGAN 2D using CNN model on the beams2d problem. `--track` will track the run on wandb, `--wandb-entity None` will use the default wandb entity, `--save-model` will save the model, `--n-epochs 200` will run for 200 epochs, and `--seed 1` will set the random seed.
 
+For reproducible debugging runs, you can additionally enable strict deterministic mode:
+```
+python engiopt/cgan_cnn_2d/cgan_cnn_2d.py --problem-id "beams2d" --seed 1 --strict-determinism
+```
+This enables stricter PyTorch deterministic settings and deterministic data shuffling while keeping the default behavior unchanged when the flag is omitted.
+
 You can always check the help for more options:
 ```
 python engiopt/cgan_cnn_2d/cgan_cnn_2d.py -h

engiopt/cgan_1d/cgan_1d.py

@@ -8,7 +8,6 @@
 
 from dataclasses import dataclass
 import os
-import random
 import time
 from typing import TYPE_CHECKING
 
@@ -21,9 +20,12 @@
 from torchvision import transforms
 import tqdm
 import tyro
+import wandb
 
+from engiopt.reproducibility import enable_strict_determinism
+from engiopt.reproducibility import make_dataloader_generator
+from engiopt.reproducibility import seed_training
 from engiopt.transforms import flatten_dict_factory
-import wandb
 
 if TYPE_CHECKING:
     from engibench.utils.problem import Problem
@@ -64,6 +66,9 @@ class Args:
     """Wandb entity name."""
     seed: int = 1
     """Random seed."""
+
+    strict_determinism: bool = False
+    """Enable strict deterministic operations for reproducibility debugging."""
     save_model: bool = False
     """Saves the model to disk."""
 
@@ -220,10 +225,9 @@ def prepare_data(problem: Problem, device: th.device) -> tuple[th.utils.data.Ten
         wandb.init(project=args.wandb_project, entity=args.wandb_entity, config=vars(args), save_code=True, name=run_name)
 
     # Seeding
-    th.manual_seed(args.seed)
-    rng = np.random.default_rng(args.seed)
-    random.seed(args.seed)
-    th.backends.cudnn.deterministic = True
+    rng = seed_training(args.seed)
+    if args.strict_determinism:
+        enable_strict_determinism(warn_only=True)
 
     os.makedirs("images", exist_ok=True)
 
@@ -240,6 +244,7 @@ def prepare_data(problem: Problem, device: th.device) -> tuple[th.utils.data.Ten
         training_ds,
         batch_size=args.batch_size,
         shuffle=True,
+        generator=make_dataloader_generator(args.seed),
     )
 
     # Loss function

engiopt/cgan_2d/cgan_2d.py

@@ -8,7 +8,6 @@
 
 from dataclasses import dataclass
 import os
-import random
 import time
 
 from engibench.utils.all_problems import BUILTIN_PROBLEMS
@@ -18,9 +17,12 @@
 from torch import nn
 import tqdm
 import tyro
-
 import wandb
 
+from engiopt.reproducibility import enable_strict_determinism
+from engiopt.reproducibility import make_dataloader_generator
+from engiopt.reproducibility import seed_training
+
 
 @dataclass
 class Args:
@@ -40,6 +42,9 @@ class Args:
     """Wandb entity name."""
     seed: int = 1
     """Random seed."""
+
+    strict_determinism: bool = False
+    """Enable strict deterministic operations for reproducibility debugging."""
     save_model: bool = False
     """Saves the model to disk."""
 
@@ -147,10 +152,9 @@ def forward(self, design: th.Tensor, conds: th.Tensor) -> th.Tensor:
         wandb.init(project=args.wandb_project, entity=args.wandb_entity, config=vars(args), save_code=True, name=run_name)
 
     # Seeding
-    th.manual_seed(args.seed)
-    rng = np.random.default_rng(args.seed)
-    random.seed(args.seed)
-    th.backends.cudnn.deterministic = True
+    rng = seed_training(args.seed)
+    if args.strict_determinism:
+        enable_strict_determinism(warn_only=True)
 
     os.makedirs("images", exist_ok=True)
 
@@ -182,6 +186,7 @@ def forward(self, design: th.Tensor, conds: th.Tensor) -> th.Tensor:
         training_ds,
         batch_size=args.batch_size,
         shuffle=True,
+        generator=make_dataloader_generator(args.seed),
     )
 
     # Optimizers

engiopt/cgan_bezier/cgan_bezier.py

@@ -8,7 +8,6 @@
 
 from dataclasses import dataclass
 import os
-import random
 import time
 from typing import TYPE_CHECKING
 
@@ -19,9 +18,12 @@
 from torch import nn
 import torch.nn.functional as f
 import tyro
-
 import wandb
 
+from engiopt.reproducibility import enable_strict_determinism
+from engiopt.reproducibility import make_dataloader_generator
+from engiopt.reproducibility import seed_training
+
 if TYPE_CHECKING:
     from collections.abc import Callable
 
@@ -46,6 +48,9 @@ class Args:
     """Wandb entity name."""
     seed: int = 6
     """Random seed."""
+
+    strict_determinism: bool = False
+    """Enable strict deterministic operations for reproducibility debugging."""
     save_model: bool = True
     """Saves the model to disk."""
 
@@ -438,11 +443,9 @@ def denormalize(self, x: th.Tensor) -> th.Tensor:
             save_code=True,
             name=run_name,
         )
-
-    th.manual_seed(args.seed)
-    rng = np.random.default_rng(args.seed)
-    random.seed(args.seed)
-    th.backends.cudnn.deterministic = True
+    rng = seed_training(args.seed)
+    if args.strict_determinism:
+        enable_strict_determinism(warn_only=True)
 
     os.makedirs("images", exist_ok=True)
 
@@ -476,6 +479,7 @@ def denormalize(self, x: th.Tensor) -> th.Tensor:
         training_ds,
         batch_size=args.batch_size,
         shuffle=True,
+        generator=make_dataloader_generator(args.seed),
     )
 
     discriminator = Discriminator(

engiopt/cgan_cnn_2d/cgan_cnn_2d.py

@@ -7,20 +7,21 @@
 
 from dataclasses import dataclass
 import os
-import random
 import time
 
 from engibench.utils.all_problems import BUILTIN_PROBLEMS
 import matplotlib.pyplot as plt
-import numpy as np
 import torch as th
 from torch import nn
 from torchvision import transforms
 import tqdm
 import tyro
-
 import wandb
 
+from engiopt.reproducibility import enable_strict_determinism
+from engiopt.reproducibility import make_dataloader_generator
+from engiopt.reproducibility import seed_training
+
 
 @dataclass
 class Args:
@@ -40,6 +41,9 @@ class Args:
     """Wandb entity name."""
     seed: int = 1
     """Random seed."""
+
+    strict_determinism: bool = False
+    """Enable strict deterministic operations for reproducibility debugging."""
     save_model: bool = False
     """Saves the model to disk."""
 
@@ -243,10 +247,9 @@ def forward(self, x: th.Tensor, c: th.Tensor) -> th.Tensor:
         wandb.init(project=args.wandb_project, entity=args.wandb_entity, config=vars(args), save_code=True, name=run_name)
 
     # Seeding
-    th.manual_seed(args.seed)
-    rng = np.random.default_rng(args.seed)
-    random.seed(args.seed)
-    th.backends.cudnn.deterministic = True
+    rng = seed_training(args.seed)
+    if args.strict_determinism:
+        enable_strict_determinism(warn_only=True)
 
     os.makedirs("images", exist_ok=True)
 
@@ -277,6 +280,7 @@ def forward(self, x: th.Tensor, c: th.Tensor) -> th.Tensor:
         training_ds,
         batch_size=args.batch_size,
         shuffle=True,
+        generator=make_dataloader_generator(args.seed),
     )
 
     # Optimizers

engiopt/cgan_cnn_3d/cgan_cnn_3d.py

@@ -9,7 +9,6 @@
 
 from dataclasses import dataclass
 import os
-import random
 import time
 
 from engibench.utils.all_problems import BUILTIN_PROBLEMS
@@ -21,10 +20,13 @@
 from torch.nn import functional
 import tqdm
 import tyro
+import wandb
 
 from engiopt.metrics import dpp_diversity
 from engiopt.metrics import mmd
-import wandb
+from engiopt.reproducibility import enable_strict_determinism
+from engiopt.reproducibility import make_dataloader_generator
+from engiopt.reproducibility import seed_training
 
 
 @dataclass
@@ -47,6 +49,9 @@ class Args:
     """Wandb entity name."""
     seed: int = 1
     """Random seed."""
+
+    strict_determinism: bool = False
+    """Enable strict deterministic operations for reproducibility debugging."""
     save_model: bool = False
     """Saves the model to disk."""
 
@@ -368,10 +373,9 @@ def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, d
         wandb.init(project=args.wandb_project, entity=args.wandb_entity, config=vars(args), save_code=True, name=run_name)
 
     # Seeding
-    th.manual_seed(args.seed)
-    rng = np.random.default_rng(args.seed)
-    random.seed(args.seed)
-    th.backends.cudnn.deterministic = True
+    rng = seed_training(args.seed)
+    if args.strict_determinism:
+        enable_strict_determinism(warn_only=True)
 
     os.makedirs("images_3d", exist_ok=True)
 
@@ -416,6 +420,7 @@ def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, d
         training_ds,
         batch_size=args.batch_size,
         shuffle=True,
+        generator=make_dataloader_generator(args.seed),
     )
 
     # Optimizers

engiopt/cgan_vae/cgan_vae.py

@@ -9,7 +9,6 @@
 
 from dataclasses import dataclass
 import os
-import random
 import time
 
 from engibench.utils.all_problems import BUILTIN_PROBLEMS
@@ -21,10 +20,13 @@
 from torch.nn import functional
 import tqdm
 import tyro
+import wandb
 
 from engiopt.metrics import dpp_diversity
 from engiopt.metrics import mmd
-import wandb
+from engiopt.reproducibility import enable_strict_determinism
+from engiopt.reproducibility import make_dataloader_generator
+from engiopt.reproducibility import seed_training
 
 
 @dataclass
@@ -47,6 +49,9 @@ class Args:
     """Wandb entity name."""
     seed: int = 1
     """Random seed."""
+
+    strict_determinism: bool = False
+    """Enable strict deterministic operations for reproducibility debugging."""
     save_model: bool = False
     """Saves the model to disk."""
 
@@ -461,10 +466,9 @@ def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, d
         wandb.init(project=args.wandb_project, entity=args.wandb_entity, config=vars(args), save_code=True, name=run_name)
 
     # Seeding
-    th.manual_seed(args.seed)
-    rng = np.random.default_rng(args.seed)
-    random.seed(args.seed)
-    th.backends.cudnn.deterministic = True
+    rng = seed_training(args.seed)
+    if args.strict_determinism:
+        enable_strict_determinism(warn_only=True)
 
     os.makedirs("images_3d", exist_ok=True)
 
@@ -517,6 +521,7 @@ def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, d
         training_ds,
         batch_size=args.batch_size,
         shuffle=True,
+        generator=make_dataloader_generator(args.seed),
     )
 
     # Optimizers for all components

engiopt/diffusion_1d/diffusion_1d.py

@@ -7,7 +7,6 @@
 
 from dataclasses import dataclass
 import os
-import random
 import time
 from typing import TYPE_CHECKING
 
@@ -21,9 +20,12 @@
 from torchvision import transforms
 import tqdm
 import tyro
+import wandb
 
+from engiopt.reproducibility import enable_strict_determinism
+from engiopt.reproducibility import make_dataloader_generator
+from engiopt.reproducibility import seed_training
 from engiopt.transforms import flatten_dict_factory
-import wandb
 
 if TYPE_CHECKING:
     from engibench.utils.problem import Problem
@@ -100,6 +102,9 @@ class Args:
     """Wandb entity name."""
     seed: int = 1
     """Random seed."""
+
+    strict_determinism: bool = False
+    """Enable strict deterministic operations for reproducibility debugging."""
     save_model: bool = False
     """Saves the model to disk."""
 
@@ -153,10 +158,9 @@ class Args:
         wandb.init(project=args.wandb_project, entity=args.wandb_entity, config=vars(args), save_code=True, name=run_name)
 
     # Seeding
-    th.manual_seed(args.seed)
-    rng = np.random.default_rng(args.seed)
-    random.seed(args.seed)
-    th.backends.cudnn.deterministic = True
+    rng = seed_training(args.seed)
+    if args.strict_determinism:
+        enable_strict_determinism(warn_only=True)
 
     os.makedirs("images", exist_ok=True)
 
@@ -188,6 +192,7 @@ class Args:
         training_ds,
         batch_size=args.batch_size,
         shuffle=True,
+        generator=make_dataloader_generator(args.seed),
     )
 
     # Optimizers