| """Deterministic seeding across python / numpy / torch.""" | |
| from __future__ import annotations | |
| import os | |
| import random | |
| def set_seed(seed: int, deterministic: bool = True) -> None: | |
| """Seed all RNGs used in training/eval. | |
| Mirrors the upstream Transolver/Geo-FNO seeding (``torch.manual_seed``, | |
| ``np.random.seed``, cudnn deterministic) and extends it to python's ``random`` | |
| and the ``PYTHONHASHSEED`` env var so runs are reproducible across seeds {0,1,2}. | |
| """ | |
| import numpy as np | |
| import torch | |
| os.environ["PYTHONHASHSEED"] = str(seed) | |
| random.seed(seed) | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| if torch.cuda.is_available(): | |
| torch.cuda.manual_seed_all(seed) | |
| if deterministic: | |
| torch.backends.cudnn.deterministic = True | |
| torch.backends.cudnn.benchmark = False | |
| # NOTE: torch.use_deterministic_algorithms(True) was tried but it coincided with a worse | |
| # Transolver-baseline result on GPU (baseline eager ~0.0090 vs ~0.0068 without it); we did | |
| # not isolate the cause and do not force it. Reproducibility is handled by averaging seeds | |
| # and disclosing run-to-run variance (README + model card). See PART 6 caveat 6. | |