| 6 | |
| 7 | |
| 8 | class RandomState: |
| 9 | def __init__(self): |
| 10 | self.random_mod_state = random.getstate() |
| 11 | self.np_state = np.random.get_state() |
| 12 | self.torch_cpu_state = torch.get_rng_state() |
| 13 | self.torch_gpu_states = [ |
| 14 | torch.cuda.get_rng_state(d) |
| 15 | for d in range(torch.cuda.device_count()) |
| 16 | ] |
| 17 | |
| 18 | def restore(self): |
| 19 | random.setstate(self.random_mod_state) |
| 20 | np.random.set_state(self.np_state) |
| 21 | torch.set_rng_state(self.torch_cpu_state) |
| 22 | for d, state in enumerate(self.torch_gpu_states): |
| 23 | torch.cuda.set_rng_state(state, d) |
| 24 | |
| 25 | |
| 26 | class RandomContext: |