77
88import ignite .distributed as idist
99import torch
10- import yaml
1110from ignite .contrib .engines import common
1211from ignite .engine import Engine
1312
3534
3635#::: } :::#
3736from ignite .utils import setup_logger
37+ from omegaconf import DictConfig , OmegaConf
3838
3939
4040def get_default_parser ():
@@ -57,17 +57,11 @@ def setup_config(parser=None):
5757 args = parser .parse_args ()
5858 config_path = args .config
5959
60- with open (config_path , "r" ) as f :
61- config = yaml .safe_load (f .read ())
60+ config = OmegaConf .load (config_path )
6261
63- optional_attributes = ["train_epoch_length" , "eval_epoch_length" ]
64- for attr in optional_attributes :
65- config [attr ] = config .get (attr , None )
62+ config .backend = args .backend
6663
67- for k , v in config .items ():
68- setattr (args , k , v )
69-
70- return args
64+ return DictConfig (config )
7165
7266
7367def log_metrics (engine : Engine , tag : str ) -> None :
@@ -138,6 +132,12 @@ def setup_output_dir(config: Any, rank: int) -> Path:
138132 return Path (idist .broadcast (config .output_dir , src = 0 ))
139133
140134
135+ def save_config (config , output_dir ):
136+ """Save configuration to config-lock.yaml for result reproducibility."""
137+ with open (f"{ output_dir } /config-lock.yaml" , "w" ) as f :
138+ OmegaConf .save (config , f )
139+
140+
141141def setup_logging (config : Any ) -> Logger :
142142 """Setup logger with `ignite.utils.setup_logger()`.
143143
0 commit comments