summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-03-04 13:20:52 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-03-04 13:20:52 +0800
commit96c1f9fa2fc747ff3c54c7f06e65706ce27fccfd (patch)
tree25e301c17725407ee143bb0d04119fe1332a87f8
parent0527c5b657c7b4fdfd7d57bf9bc5334eac480731 (diff)
Set seed for reproducibility
-rw-r--r--models/model.py11
1 files changed, 7 insertions, 4 deletions
diff --git a/models/model.py b/models/model.py
index acccbff..83b970a 100644
--- a/models/model.py
+++ b/models/model.py
@@ -1,4 +1,5 @@
import os
+import random
from datetime import datetime
from typing import Union, Optional
@@ -199,14 +200,17 @@ class Model:
self.writer = SummaryWriter(self._log_name)
+ # Set seeds for reproducibility
+ random.seed(0)
+ torch.manual_seed(0)
self.rgb_pn.train()
# Init weights at first iter
if self.curr_iter == 0:
self.rgb_pn.apply(self.init_weights)
else: # Load saved state dicts
checkpoint = torch.load(self._checkpoint_name)
- iter_, loss = checkpoint['iter'], checkpoint['loss']
- print('{0:5d} loss: {1:.3f}'.format(iter_, loss))
+ random.setstate(checkpoint['rand_states'][0])
+ torch.set_rng_state(checkpoint['rand_states'][1])
self.rgb_pn.load_state_dict(checkpoint['model_state_dict'])
self.optimizer.load_state_dict(checkpoint['optim_state_dict'])
self.scheduler.load_state_dict(checkpoint['sched_state_dict'])
@@ -320,11 +324,10 @@ class Model:
if self.curr_iter % 1000 == 0:
torch.save({
- 'iter': self.curr_iter,
+ 'rand_states': (random.getstate(), torch.get_rng_state()),
'model_state_dict': self.rgb_pn.state_dict(),
'optim_state_dict': self.optimizer.state_dict(),
'sched_state_dict': self.scheduler.state_dict(),
- 'loss': loss,
}, self._checkpoint_name)
if self.curr_iter == self.total_iter: