Change default Namespace to Omegaconf and new method to save config-lock.yaml for reproducibility (#292)

* Added omegaconf as default Namespace for config and tests for saving the config

* tests fix

* remove shutil.copy from templates

* Change the save_config function and added the None parameters in config.yaml

* remove yaml import and replace it with Omegaconf methods

* fix yaml import in test_all.py

Author: guptaaryan16

scripts/requirements.txt

@@ -8,3 +8,4 @@ pytest
 transformers
 datasets
 tensorboard
+omegaconf

src/templates/template-common/config.yaml

@@ -6,6 +6,8 @@ num_workers: 4
 max_epochs: 20
 use_amp: false
 debug: false
+train_epoch_length: null
+eval_epoch_length: null
 
 #::: if (it.dist === 'spawn') { :::#
 # distributed spawn

src/templates/template-common/requirements.txt

@@ -2,6 +2,7 @@ torch>=1.10.2
 torchvision>=0.11.3
 pytorch-ignite>=0.4.8
 pyyaml
+omegaconf
 
 #::: if (['neptune', 'polyaxon'].includes(it.logger)) { :::#
 

src/templates/template-common/test_all.py

@@ -0,0 +1,10 @@
+def test_save_config():
+    with open("./config.yaml", "r") as f:
+        config = OmegaConf.load(f)
+
+    save_config(config, "./")
+
+    with open("./config-lock.yaml", "r") as f:
+        test_config = OmegaConf.load(f)
+
+    assert config == test_config

src/templates/template-common/utils.py

@@ -7,7 +7,6 @@
 
 import ignite.distributed as idist
 import torch
-import yaml
 from ignite.contrib.engines import common
 from ignite.engine import Engine
 
@@ -35,6 +34,7 @@
 
 #::: } :::#
 from ignite.utils import setup_logger
+from omegaconf import DictConfig, OmegaConf
 
 
 def get_default_parser():
@@ -57,17 +57,11 @@ def setup_config(parser=None):
     args = parser.parse_args()
     config_path = args.config
 
-    with open(config_path, "r") as f:
-        config = yaml.safe_load(f.read())
+    config = OmegaConf.load(config_path)
 
-    optional_attributes = ["train_epoch_length", "eval_epoch_length"]
-    for attr in optional_attributes:
-        config[attr] = config.get(attr, None)
+    config.backend = args.backend
 
-    for k, v in config.items():
-        setattr(args, k, v)
-
-    return args
+    return DictConfig(config)
 
 
 def log_metrics(engine: Engine, tag: str) -> None:
@@ -138,6 +132,12 @@ def setup_output_dir(config: Any, rank: int) -> Path:
     return Path(idist.broadcast(config.output_dir, src=0))
 
 
+def save_config(config, output_dir):
+    """Save configuration to config-lock.yaml for result reproducibility."""
+    with open(f"{output_dir}/config-lock.yaml", "w") as f:
+        OmegaConf.save(config, f)
+
+
 def setup_logging(config: Any) -> Logger:
     """Setup logger with `ignite.utils.setup_logger()`.
 

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

@@ -1,6 +1,5 @@
 import os
 from pprint import pformat
-from shutil import copy
 from typing import Any
 
 import ignite.distributed as idist
@@ -25,7 +24,7 @@ def run(local_rank: int, config: Any):
     # create output folder and copy config file to output dir
     config.output_dir = setup_output_dir(config, rank)
     if rank == 0:
-        copy(config.config, f"{config.output_dir}/config-lock.yaml")
+        save_config(config, config.output_dir)
 
     # donwload datasets and create dataloaders
     dataloader_train, dataloader_eval = setup_data(config)
@@ -69,7 +68,7 @@ def run(local_rank: int, config: Any):
     # setup engines logger with python logging
     # print training configurations
     logger = setup_logging(config)
-    logger.info("Configuration: \n%s", pformat(vars(config)))
+    logger.info("Configuration: \n%s", pformat(config))
     trainer.logger = evaluator.logger = logger
 
     trainer.add_event_handler(Events.ITERATION_COMPLETED, lr_scheduler)

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

@@ -5,9 +5,11 @@
 import ignite.distributed as idist
 import pytest
 from data import setup_data
+from omegaconf import OmegaConf
 from torch import nn, optim
 from torch.functional import Tensor
 from torch.utils.data import DataLoader
+from utils import save_config
 
 
 def set_up():
@@ -45,3 +47,6 @@ def test_setup_data():
     assert isinstance(eval_batch["attention_mask"], Tensor)
     assert isinstance(eval_batch["token_type_ids"], Tensor)
     assert isinstance(eval_batch["label"], Tensor)
+
+
+#::= from_template_common ::#

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

@@ -1,5 +1,4 @@
 from pprint import pformat
-from shutil import copy
 from typing import Any
 
 import ignite.distributed as idist
@@ -22,7 +21,7 @@ def run(local_rank: int, config: Any):
     # create output folder and copy config file to output dir
     config.output_dir = setup_output_dir(config, rank)
     if rank == 0:
-        copy(config.config, f"{config.output_dir}/config-lock.yaml")
+        save_config(config, config.output_dir)
 
     # donwload datasets and create dataloaders
     dataloader_train, dataloader_eval = setup_data(config)
@@ -59,7 +58,7 @@ def run(local_rank: int, config: Any):
     # setup engines logger with python logging
     # print training configurations
     logger = setup_logging(config)
-    logger.info("Configuration: \n%s", pformat(vars(config)))
+    logger.info("Configuration: \n%s", pformat(config))
     trainer.logger = evaluator.logger = logger
 
     trainer.add_event_handler(Events.ITERATION_COMPLETED, lr_scheduler)

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

@@ -6,9 +6,11 @@
 import pytest
 import torch
 from data import setup_data
+from omegaconf import OmegaConf
 from torch import nn, optim, Tensor
 from torch.utils.data.dataloader import DataLoader
 from trainers import setup_evaluator
+from utils import save_config
 
 
 def set_up():
@@ -48,3 +50,6 @@ def test_setup_evaluator():
     evaluator = setup_evaluator(config, model, device)
     evaluator.run([batch, batch])
     assert isinstance(evaluator.state.output, tuple)
+
+
+#::= from_template_common ::#

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

@@ -1,5 +1,4 @@
 from pprint import pformat
-from shutil import copy
 from typing import Any
 
 import ignite.distributed as idist
@@ -25,7 +24,7 @@ def run(local_rank: int, config: Any):
     # create output folder and copy config file to output dir
     config.output_dir = setup_output_dir(config, rank)
     if rank == 0:
-        copy(config.config, f"{config.output_dir}/config-lock.yaml")
+        save_config(config, config.output_dir)
 
     # donwload datasets and create dataloaders
     dataloader_train, dataloader_eval, num_channels = setup_data(config)
@@ -74,7 +73,7 @@ def run(local_rank: int, config: Any):
     # setup engines logger with python logging
     # print training configurations
     logger = setup_logging(config)
-    logger.info("Configuration: \n%s", pformat(vars(config)))
+    logger.info("Configuration: \n%s", pformat(config))
     trainer.logger = evaluator.logger = logger
 
     #::: if (it.save_training || it.save_evaluation) { :::#