File size: 6,162 Bytes
256a159 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
import inspect
from abc import abstractmethod
from copy import deepcopy
from typing import Dict, List, Optional
from mmengine.config import ConfigDict
from opencompass.utils import (dataset_abbr_from_cfg, get_logger,
model_abbr_from_cfg, task_abbr_from_cfg)
class BasePartitioner:
"""Base class for partitioners. A partitioner is responsible for
partitioning the config into tasks.
Args:
out_dir (str): The output directory of tasks.
keep_keys (Optional[List[str]], optional): The keys to be kept from the
experiment config to the task config. Defaults to None. If None,
the following keys will be kept:
- eval.runner.task.judge_cfg
- eval.runner.task.dump_details
"""
def __init__(self, out_dir: str, keep_keys: Optional[List[str]] = None):
self.logger = get_logger()
self.out_dir = out_dir
if keep_keys is None:
self.keep_keys = [
'eval.runner.task.judge_cfg',
'eval.runner.task.dump_details',
]
else:
self.keep_keys = keep_keys
def __call__(self, cfg: ConfigDict) -> List[Dict]:
"""Generate tasks from config. Each task is defined as a
dict and will run independently as a unit. Its structure is as
follows:
.. code-block:: python
{
'models': [], # a list of model configs
'datasets': [[]], # a nested list of dataset configs, each
list corresponds to a model
'work_dir': '', # the work dir
}
Args:
cfg (ConfigDict): The config dict, containing "models", "dataset"
and "work_dir" keys.
Returns:
List[Dict]: A list of tasks.
"""
cfg = deepcopy(cfg)
work_dir = cfg['work_dir']
add_cfg = {}
for k in self.keep_keys:
try:
key_chain = k.split('.')
ori_ptr = cfg
tgt_ptr = add_cfg
for key in key_chain[:-1]:
ori_ptr = ori_ptr[key]
if key not in tgt_ptr:
tgt_ptr[key] = {}
tgt_ptr = tgt_ptr[key]
tgt_ptr[key_chain[-1]] = ori_ptr[key_chain[-1]]
except Exception:
self.logger.debug(f'Key {k} not found in config, ignored.')
self.logger.debug(f'Additional config: {add_cfg}')
model_and_dataset_args = self.parse_model_dataset_args(cfg)
tasks = self.partition(**model_and_dataset_args,
work_dir=work_dir,
out_dir=self.out_dir,
add_cfg=add_cfg)
self.logger.info(f'Partitioned into {len(tasks)} tasks.')
for i, task in enumerate(tasks):
self.logger.debug(f'Task {i}: {task_abbr_from_cfg(task)}')
return tasks
def parse_model_dataset_args(self, cfg: ConfigDict):
models = cfg['models']
datasets = cfg['datasets']
sig = inspect.signature(self.partition)
if 'model_dataset_combinations' in sig.parameters:
combs = cfg.get('model_dataset_combinations', None)
if combs is None:
combs = [{'models': models, 'datasets': datasets}]
else:
# sanity check
model_abbrs = [model_abbr_from_cfg(model) for model in models]
dataset_abbrs = [
dataset_abbr_from_cfg(dataset) for dataset in datasets
]
for comb in combs:
for model in comb['models']:
if model_abbr_from_cfg(model) not in model_abbrs:
raise ValueError(
f'Model {model_abbr_from_cfg(model)} '
'not found in config.')
for dataset in comb['datasets']:
if dataset_abbr_from_cfg(dataset) not in dataset_abbrs:
raise ValueError(
f'Dataset {dataset_abbr_from_cfg(dataset)} '
'not found in config.')
used_kwargs = {'model_dataset_combinations': combs}
else:
if cfg.get('model_dataset_combinations', None) is not None:
self.logger.warning(
'model_dataset_combinations is not supported by '
f'{self.__class__.__name__}. Ignored.')
used_kwargs = {'models': models, 'datasets': datasets}
return used_kwargs
@abstractmethod
def partition(self,
models: List[ConfigDict],
datasets: List[ConfigDict],
work_dir: str,
out_dir: str,
add_cfg: Dict = {}) -> List[Dict]:
"""Partition model-dataset pairs into tasks. Each task is defined as a
dict and will run independently as a unit. Its structure is as
follows:
.. code-block:: python
{
'models': [], # a list of model configs
'datasets': [[]], # a nested list of dataset configs, each
list corresponds to a model
'work_dir': '', # the work dir
**add_cfg # other keys to be added in the config
}
Args:
models (List[ConfigDict]): A list of model configs.
datasets (List[ConfigDict]): A list of dataset configs.
work_dir (str): The work dir for the task.
out_dir (str): The full output path for the task, intended for
Partitioners to check whether the task is finished via the
existency of result file in this directory.
add_cfg (dict): Other common keys to be added in the task config,
used to share the same config among tasks. Defaults to {}.
Returns:
List[Dict]: A list of tasks.
"""
|