Merge pull request #106 from sarvghotra/bug_fix_dist_sampler

Bug fix in dist sampler that caused same data order in each epoch
cszn · Apr 21, 2022 · f4573f3 · f4573f3
2 parents 7e51c16 + fc2f79e
commit f4573f3
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/main_train_drunet.py b/main_train_drunet.py
@@ -51,7 +51,7 @@ def main(json_path='options/train_drunet.json'):
     if opt['dist']:
         init_dist('pytorch')
     opt['rank'], opt['world_size'] = get_dist_info()
-    
+
     if opt['rank'] == 0:
         util.mkdirs((path for key, path in opt['path'].items() if 'pretrained' not in key))
 
@@ -160,6 +160,9 @@ def main(json_path='options/train_drunet.json'):
     '''
 
     for epoch in range(1000000):  # keep running
+        if opt['dist']:
+            train_sampler.set_epoch(epoch)
+
         for i, train_data in enumerate(train_loader):
 
             current_step += 1

diff --git a/main_train_gan.py b/main_train_gan.py
@@ -172,6 +172,9 @@ def main(json_path='options/train_msrresnet_gan.json'):
     '''
 
     for epoch in range(1000000):  # keep running
+        if opt['dist']:
+            train_sampler.set_epoch(epoch)
+
         for i, train_data in enumerate(train_loader):
 
             current_step += 1

diff --git a/main_train_psnr.py b/main_train_psnr.py
@@ -54,7 +54,7 @@ def main(json_path='options/train_msrresnet_psnr.json'):
     if opt['dist']:
         init_dist('pytorch')
     opt['rank'], opt['world_size'] = get_dist_info()
-    
+
     if opt['rank'] == 0:
         util.mkdirs((path for key, path in opt['path'].items() if 'pretrained' not in key))
 
@@ -165,6 +165,9 @@ def main(json_path='options/train_msrresnet_psnr.json'):
     '''
 
     for epoch in range(1000000):  # keep running
+        if opt['dist']:
+            train_sampler.set_epoch(epoch)
+
         for i, train_data in enumerate(train_loader):
 
             current_step += 1