-
-
Save maulikmadhavi/68b7d90d15cff959928732d25ef79a87 to your computer and use it in GitHub Desktop.
pytorch - set seed everything
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # Tested regorously on multiple python environment and multiple devices | |
| # Feel free to update | |
| import numpy as np | |
| from torch.utils.data import Dataset, DataLoader | |
| import random | |
| import torch | |
| import os | |
| from functools import partial | |
| print("Numpy version: ", np.__version__) | |
| print("torch version: ", torch.__version__) | |
| class Fix_seed(object): | |
| def __init__(self, seed): | |
| self.seed = seed | |
| self.gen = self.seed_everything() | |
| self.DataLoader_ = partial(DataLoader, | |
| generator=self.gen, | |
| worker_init_fn=self.worker_init_fn | |
| ) | |
| def seed_everything(self): | |
| print("seed everythin") | |
| seed = self.seed | |
| random.seed(seed) | |
| os.environ['PYTHONHASHSEED'] = str(seed) | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| torch.cuda.manual_seed(seed) | |
| torch.backends.cudnn.deterministic = True | |
| @staticmethod | |
| def worker_init_fn(worker_id): | |
| worker_seed = torch.initial_seed() % 2**32 | |
| np.random.seed(worker_seed) | |
| random.seed(worker_seed) | |
| class RandomDataset(Dataset): | |
| def __init__(self) -> None: | |
| super().__init__() | |
| self.data = np.arange(0,16).reshape(-1,2) | |
| def __getitem__(self, index): | |
| # return np.random.randint(0, 1000, 3) # 3 random numbers size: (3,) | |
| # print(random.random()) | |
| return self.data[index], random.random() # image size: (2, 16) | |
| def __len__(self): | |
| return self.data.shape[0] | |
| dataset = RandomDataset() | |
| seeder = Fix_seed(1234) | |
| dataloader = seeder.DataLoader_(dataset, batch_size=2, num_workers=16, shuffle=True) | |
| EPOCHS = 3 | |
| for _ in range(EPOCHS): | |
| print(f"================epoch: {_} ====================") | |
| for batch in dataloader: | |
| print(batch) | |
| # print(batch.shape) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment