python用yaml装参数并支持命令行修改

效果：

• 将实验用的参数写入 yaml 文件，而不是全部用 argparse 传，否则命令会很长；
• 同时支持在命令行临时加、改一些参数，避免事必要在 yaml 中改参数，比较灵活（如 grid-search 时遍历不同的 loss weights）。

最初是在 MMDetection 中看到这种写法，参考 [1] 中 --cfg-options 这个参数，核心是 DictAction 类，定义在 [2]。yaml 一些支持的写法参考 [3]。本文同时作为 python yaml 读、写简例。

Code

• DictAction 类抄自 [2]；
• parse_cfg 函数读 yaml 参数，并按命令行输入加、改参数（覆盖 yaml），用 EasyDict 装；
• 用 yaml 备份参数时，用 easydict2dict 将 EasyDict 递归改回 dict，yaml 会干净点。不用也行。

from argparse import Action, ArgumentParser, Namespace
import copy
from easydict import EasyDict
from typing import Any, Optional, Sequence, Tuple, Union
import yamlclass DictAction(Action):"""抄自 MMEngineargparse action to split an argument into KEY=VALUE formon the first = and append to a dictionary. List options canbe passed as comma separated values, i.e 'KEY=V1,V2,V3', or with explicitbrackets, i.e. 'KEY=[V1,V2,V3]'. It also support nested brackets to buildlist/tuple values. e.g. 'KEY=[(V1,V2),(V3,V4)]'"""@staticmethoddef _parse_int_float_bool(val: str) -> Union[int, float, bool, Any]:"""parse int/float/bool value in the string."""try:return int(val)except ValueError:passtry:return float(val)except ValueError:passif val.lower() in ['true', 'false']:return True if val.lower() == 'true' else Falseif val == 'None':return Nonereturn val@staticmethoddef _parse_iterable(val: str) -> Union[list, tuple, Any]:"""Parse iterable values in the string.All elements inside '()' or '[]' are treated as iterable values.Args:val (str): Value string.Returns:list | tuple | Any: The expanded list or tuple from the string,or single value if no iterable values are found.Examples:>>> DictAction._parse_iterable('1,2,3')[1, 2, 3]>>> DictAction._parse_iterable('[a, b, c]')['a', 'b', 'c']>>> DictAction._parse_iterable('[(1, 2, 3), [a, b], c]')[(1, 2, 3), ['a', 'b'], 'c']"""def find_next_comma(string):"""Find the position of next comma in the string.If no ',' is found in the string, return the string length. Allchars inside '()' and '[]' are treated as one element and thus ','inside these brackets are ignored."""assert (string.count('(') == string.count(')')) and (string.count('[') == string.count(']')), \f'Imbalanced brackets exist in {string}'end = len(string)for idx, char in enumerate(string):pre = string[:idx]# The string before this ',' is balancedif ((char == ',') and (pre.count('(') == pre.count(')'))and (pre.count('[') == pre.count(']'))):end = idxbreakreturn end# Strip ' and " characters and replace whitespace.val = val.strip('\'\"').replace(' ', '')is_tuple = Falseif val.startswith('(') and val.endswith(')'):is_tuple = Trueval = val[1:-1]elif val.startswith('[') and val.endswith(']'):val = val[1:-1]elif ',' not in val:# val is a single valuereturn DictAction._parse_int_float_bool(val)values = []while len(val) > 0:comma_idx = find_next_comma(val)element = DictAction._parse_iterable(val[:comma_idx])values.append(element)val = val[comma_idx + 1:]if is_tuple:return tuple(values)return valuesdef __call__(self,parser: ArgumentParser,namespace: Namespace,values: Union[str, Sequence[Any], None],option_string: str = None):"""Parse Variables in string and add them into argparser.Args:parser (ArgumentParser): Argument parser.namespace (Namespace): Argument namespace.values (Union[str, Sequence[Any], None]): Argument string.option_string (list[str], optional): Option string.Defaults to None."""# Copied behavior from `argparse._ExtendAction`.options = copy.copy(getattr(namespace, self.dest, None) or {})if values is not None:for kv in values:key, val = kv.split('=', maxsplit=1)options[key] = self._parse_iterable(val)setattr(namespace, self.dest, options)def parse_cfg(yaml_file, update_dict={}):"""load configurations from a yaml file & update from command-line argmentsInput:yaml_file: str, path to a yaml configuration fileupdate_dict: dict, to modify/update options in those yaml configurationsOutput:cfg: EasyDict"""with open(args.cfg, "r") as f:cfg = EasyDict(yaml.safe_load(f))if update_dict:assert isinstance(update_dict, dict)for k, v in update_dict.items():k_list = k.split('.')assert len(k_list) > 0if len(k_list) == 1: # 单级，e.g. lr=0.1cfg[k_list[0]] = velse: # 多级，e.g. optimizer.group1.lr=0.2ptr = cfgfor i, _k in enumerate(k_list):if i == len(k_list) - 1: # last layerptr[_k] = velif _k not in ptr:ptr[_k] = {}ptr = ptr[_k]return cfgdef easydict2dict(ed):"""convert EasyDict to dict for clean yaml"""d = {}for k, v in ed.items():if isinstance(v, EasyDict):d[k] = easydict2dict(v)else:d[k] = vreturn dif "__main__" == __name__:# test command:#   python config.py --cfg-options int=5 dict2.lr=8 dict2.newdict.newitem=flyimport pprintparser = ArgumentParser()parser.add_argument("--cfg", type=str, default="config.yaml", help="指定 yaml")parser.add_argument('--cfg-options',nargs='+',action=DictAction,help='override some settings in the used config, the key-value pair ''in xxx=yyy format will be merged into config file. If the value to ''be overwritten is a list, it should be like key="[a,b]" or key=a,b ''It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" ''Note that the quotation marks are necessary and that no white space ''is allowed.')args = parser.parse_args()# 命令行临时加、改参数pprint.pprint(args.cfg_options) # dict# 读 yaml，并按命令行输入加、改参数cfg = parse_cfg(args.cfg, args.cfg_options)pprint.pprint(cfg)# 备份 yaml（写 yaml）with open("backup-config.yaml", 'w') as f:yaml.dump(easydict2dict(cfg), f)

输入的 config.yaml：

• 语法参考 [3]

# An example yaml configuration file, used in utils/config.py as an example input.
# Ref: https://pyyaml.org/wiki/PyYAMLDocumentationlog_path: ./log
none: [~, null, None]
bool: [true, false, on, off, True, False]
int: 42				# <- 改它
float: 3.14159
list: [LITE, RES_ACID, SUS_DEXT]
list2:- -1- 0- 1
str:a20.2# d: tom
dict: {hp: 13, sp: 5}
dict2:				# <- 加它lr: 0.01			# <- 改它decay_rate: 0.1name: jerry

测试：

# --cfg-options 支持多级指定（用「.」分隔）
python config.py --cfg config.yaml --cfg-options int=5 dict2.lr=8 dict2.newdict.newitem=fly

输出：

{'dict2.lr': 8, 'dict2.newdict.newitem': 'fly', 'int': 5}
{'bool': [True, False, True, False, True, False],'dict': {'hp': 13, 'sp': 5},'dict2': {'decay_rate': 0.1,'lr': 8,							# <- 改了'name': 'jerry','newdict': {'newitem': 'fly'}},	# <- 加了'float': 3.14159,'int': 5,									# <- 改了'list': ['LITE', 'RES_ACID', 'SUS_DEXT'],'list2': [-1, 0, 1],'log_path': './log','none': [None, None, 'None'],'str': 'a 2 0.2'}

References

1. open-mmlab/mmdetection/tools/train.py
2. open-mmlab/mmengine/mmengine/config/config.py
3. PyYAML Documentation