Fix the output_dir logic and bugs for reproducibility (#307)

* Changed the function setup_output_dir and output_dir logic in templates

* add coment again to setup_output_dir

* Modify test_save_config for the new behaviour

* Add temp directory for test_save_config

Author: guptaaryan16

src/templates/template-common/test_all.py

@@ -2,9 +2,15 @@ def test_save_config():
     with open("./config.yaml", "r") as f:
         config = OmegaConf.load(f)
 
-    save_config(config, "./")
+    # Add backend to config (similar to setup_config)
+    config.backend = None
 
-    with open("./config-lock.yaml", "r") as f:
-        test_config = OmegaConf.load(f)
+    with tempfile.TemporaryDirectory() as output_dir:
+        output_dir = Path(output_dir)
 
-    assert config == test_config
+        save_config(config, output_dir)
+
+        with open(output_dir / "config-lock.yaml", "r") as f:
+            test_config = OmegaConf.load(f)
+
+        assert config == test_config

src/templates/template-common/utils.py

@@ -149,14 +149,14 @@ def resume_from(
 
 def setup_output_dir(config: Any, rank: int) -> Path:
     """Create output folder."""
+    output_dir = config.output_dir
     if rank == 0:
         now = datetime.now().strftime("%Y%m%d-%H%M%S")
         name = f"{now}-backend-{config.backend}-lr-{config.lr}"
         path = Path(config.output_dir, name)
         path.mkdir(parents=True, exist_ok=True)
-        config.output_dir = path.as_posix()
-
-    return Path(idist.broadcast(config.output_dir, src=0))
+        output_dir = path.as_posix()
+    return Path(idist.broadcast(output_dir, src=0))
 
 
 def save_config(config, output_dir):

src/templates/template-text-classification/main.py

@@ -27,9 +27,11 @@ def run(local_rank: int, config: Any):
     manual_seed(config.seed + rank)
 
     # create output folder and copy config file to output dir
-    config.output_dir = setup_output_dir(config, rank)
+    output_dir = setup_output_dir(config, rank)
     if rank == 0:
-        save_config(config, config.output_dir)
+        save_config(config, output_dir)
+
+    config.output_dir = output_dir
 
     # donwload datasets and create dataloaders
     dataloader_train, dataloader_eval = setup_data(config)

src/templates/template-text-classification/test_all.py

@@ -1,5 +1,7 @@
 import os
+import tempfile
 from argparse import Namespace
+from pathlib import Path
 from typing import Iterable
 
 import ignite.distributed as idist

src/templates/template-vision-classification/main.py

@@ -24,9 +24,11 @@ def run(local_rank: int, config: Any):
     manual_seed(config.seed + rank)
 
     # create output folder and copy config file to output dir
-    config.output_dir = setup_output_dir(config, rank)
+    output_dir = setup_output_dir(config, rank)
     if rank == 0:
-        save_config(config, config.output_dir)
+        save_config(config, output_dir)
+
+    config.output_dir = output_dir
 
     # donwload datasets and create dataloaders
     dataloader_train, dataloader_eval = setup_data(config)

src/templates/template-vision-classification/test_all.py

@@ -1,5 +1,7 @@
 import os
+import tempfile
 from argparse import Namespace
+from pathlib import Path
 from typing import Iterable
 
 import ignite.distributed as idist

src/templates/template-vision-dcgan/main.py

@@ -28,9 +28,11 @@ def run(local_rank: int, config: Any):
     manual_seed(config.seed + rank)
 
     # create output folder and copy config file to output dir
-    config.output_dir = setup_output_dir(config, rank)
+    output_dir = setup_output_dir(config, rank)
     if rank == 0:
-        save_config(config, config.output_dir)
+        save_config(config, output_dir)
+
+    config.output_dir = output_dir
 
     # donwload datasets and create dataloaders
     dataloader_train, dataloader_eval, num_channels = setup_data(config)

src/templates/template-vision-dcgan/test_all.py

@@ -1,5 +1,7 @@
 import os
+import tempfile
 from argparse import Namespace
+from pathlib import Path
 from typing import Iterable
 
 import ignite.distributed as idist

src/templates/template-vision-segmentation/main.py

@@ -34,9 +34,11 @@ def run(local_rank: int, config: Any):
     manual_seed(config.seed + rank)
 
     # create output folder and copy config file to output dir
-    config.output_dir = setup_output_dir(config, rank)
+    output_dir = setup_output_dir(config, rank)
     if rank == 0:
-        save_config(config, config.output_dir)
+        save_config(config, output_dir)
+
+    config.output_dir = output_dir
 
     # donwload datasets and create dataloaders
     dataloader_train, dataloader_eval = setup_data(config)

src/templates/template-vision-segmentation/test_all.py

@@ -1,5 +1,7 @@
 import os
+import tempfile
 from argparse import Namespace
+from pathlib import Path
 
 import pytest
 from data import setup_data