fixed subsampling for ddp training

facebookresearch · Aug 13, 2020 · 560e583 · 560e583
1 parent 2b442ad
commit 560e583
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/fastmri/data/mri_data.py b/fastmri/data/mri_data.py
@@ -114,7 +114,7 @@ class SliceDataset(Dataset):
             what fraction of the volumes should be loaded.
     """
 
-    def __init__(self, root, transform, challenge, sample_rate=1):
+    def __init__(self, root, transform, challenge, sample_rate=1, seed=0):
         if challenge not in ("singlecoil", "multicoil"):
             raise ValueError('challenge should be either "singlecoil" or "multicoil"')
 
@@ -126,6 +126,7 @@ def __init__(self, root, transform, challenge, sample_rate=1):
 
         files = list(pathlib.Path(root).iterdir())
         if sample_rate < 1:
+            random.seed(seed) # get the same files in every process
             random.shuffle(files)
             num_files = round(len(files) * sample_rate)
             files = files[:num_files]