Integration between Omegaconf and argparse for mixed config file and CLI arguments
Flexible configuration management is important during experimentation, e.g. when training machine learning models.
Ideally, we want both a configuration file to hold more "stable" hyperparameter values and the flexibility to change values through command line arguments for rapid experimentation.
This package provides a barebones solution based on the excellent OmegaConf package.
Specifically, we extend the OmegaConf class with a static from_argparse method to parse arguments provided using argparse,
and provide utility functions to merge the default CLI values, YAML configuration file values and user provided CLI arguments.
Install package from PyPi:
pip install omegaconf-argparse
Then in your code you can import:
from omegacli import OmegaConf, generate_config_template, parse_config
Let's start with an example. You have an argparse based script, and you want to move to a more reproducible setup based on configuration files, without losing the flexibility of cli arguments.
In examples/mnist.py we have modified the example MNIST training script from pytorch repo.
The diff is:
11a12,13
> from omegacli import generate_config_template, parse_config
>
158a161,172
> parser.add_argument(
> "--generate-config",
> action="store_true",
> default=False,
> help="Generate example YAML configuration file",
> )
> parser.add_argument(
> "--config-path",
> type=str,
> default=None,
> help="Path to configuration file",
> )
159a174,182
>
> if args.generate_config:
> import sys
>
> generate_config_template(parser, args.config)
> sys.exit(0)
>
> args = parse_config(parser, args.config)Note we have added two command line arguments --generate-config and --config-path.
When we run
python mnist.py --generate-config --config-path config.yamlit will create a config.yaml file, which can be used from now on:
batch_size: 64
test_batch_size: 1000
epochs: null
lr: 1.0
gamma: 0.7
no_cuda: false
no_mps: false
dry_run: false
seed: 1
log_interval: 10
save_model: falseNow if we run:
python mnist.py --config config.yaml
The script will use the values provided in the config.yaml file. If we change the configuration:
lr: 0.1The training will use lr=0.1.
At any time we can run a quick experiment (let's say with gamma=1.0) and override the config values using the CLI:
python mnist.py --config config.yaml --gamma 1.0
When we are confident with our experiments we can set the best hyperparameter values in the configuration file and push to a remote repo for reproducibility.
We provide a high-level utility function parse_config that merges a YAML configuration file with an argparse.ArgumentParser:
import io
from omegacli import parse_config
mock_config_file = io.StringIO('''
model:
hidden: 100
''')
parser = argparse.ArgumentParser("My cool model")
parser.add_argument("--hidden", dest="model.hidden", type=int, default=20)
cfg = parse_config(parser, mock_config_file)
>>> {'model': {'hidden': 100}}
type(cfg)
>>> <class 'omegaconf.dictconfig.DictConfig'>
cfg = parse_config(parser, mock_config_file, args=["--hidden", "200"])
>>> {'model': {'hidden': 200}}
mock_config_file = io.StringIO('''
random_value: hello
''')
cfg = parse_config(parser, mock_config_file)
>>> {'model': {'hidden': 20}, 'random_value': 'hello'}You can also use the patched OmegaConf class directly:
import argparse
from omegacli import OmegaConf
parser = argparse.ArgumentParser("My cool model")
parser.add_argument("--hidden", dest="model.hidden", type=int, default=20)
user_provided_args, default_args = OmegaConf.from_argparse(parser, args=["--hidden", "100"])
user_provided_args
>>> {'model': {'hidden': 100}}
default_args
>>> {}
user_provided_args, default_args = OmegaConf.from_argparse(parser)
user_provided_args
>>> {}
default_args
>>> {'model': {'hidden': 20}}NOTE: the from_argparse method calls the parser.parse_args().
The precedence for merging is as follows
- user provided cli args > config file values > default cli args values
E.g.:
- if you don't include a value in your configuration it will take the default value from the argparse arguments
- if you provide a cli arg (e.g. run the script with --bsz 64) it will override the value in the config file
To create a nested configuration structure and match with the argparse provided CLI args,
we use the dest kwarg when adding an argument with parser.add_argument.
Specifically, we follow the convention that the destination is a string, delimited by . to indicate nested structure.
For example:
parser.add_argument("--hidden", dest="model_hidden", type=int, default=20)will create a top-level element in the configuration:
user_provided_args, default_args = OmegaConf.from_argparse(parser, args=["--hidden", "100"])
user_provided_args
>>> {'model_hidden': 100}This will match with the following entry in the YAML file:
model_hidden: 100The following:
parser.add_argument("--hidden", dest="model.hidden", type=int, default=20)will create a nested structure in the configuration:
user_provided_args, default_args = OmegaConf.from_argparse(parser, args=["--hidden", "100"])
user_provided_args
>>> {'model': {'hidden': 100}}and will match with the following YAML entry:
model:
hidden: 100The parsing is recursive, so you can go as deep as you want.
Run:
from omegacli import generate_config_template
generate_config_template(parser, "example-config.yaml")This will initialize a configuration file, that is consistent with the argument parser. You can use this as a starting point for saving and editing your configuration.
- Hydra: A more feature-rich, but more complex solution. If you are willing to introduce the dependency you can use it
- Pytorch-Lightning: PL introduced a similar functionality. You can use it if you are in the PL ecosystem.
OmegaConf plans to remain agnostic to the command line argument parser, therefore we cannot merge this solution upstream. See related issue.