|
from copy import deepcopy |
|
from typing import Dict, List |
|
|
|
from mmengine.config import Config, ConfigDict |
|
|
|
from opencompass.registry import PARTITIONERS |
|
|
|
from .base import BasePartitioner |
|
|
|
|
|
@PARTITIONERS.register_module() |
|
class MultimodalNaivePartitioner(BasePartitioner): |
|
"""Multimodal naive task partitioner. |
|
|
|
This partitioner will generate a task for each |
|
model-dataset-evaluator pair. |
|
|
|
Args: |
|
config (ConfigDict): The full config dict. |
|
""" |
|
|
|
def partition(self, models: List[ConfigDict], datasets: List[ConfigDict], |
|
evaluators: List[ConfigDict], load_froms: List[ConfigDict], |
|
work_dir: str, num_gpus: int, num_procs: int, |
|
launcher: str) -> 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 list of dataset configs |
|
'evaluators': [], # a list of evaluator configs |
|
'load_froms': [], # a list of load_from paths |
|
'work_dir': '', # the work dir |
|
'num_gpus': int, # integer, number of gpus for each task |
|
'num_procs': int, # integer, number of gpus on single machine |
|
'launcher': str, # string, how to launch distributed training |
|
} |
|
|
|
Args: |
|
models (List[ConfigDict]): A list of model configs. |
|
datasets (List[ConfigDict]): A list of dataset configs. |
|
evaluators (List[ConfigDict]): A list of evaluator configs. |
|
load_froms (List[ConfigDict]): A list of load_from paths. |
|
work_dir (str): The work dir for the task. |
|
num_gpus (int): Number of gpus for each task. |
|
num_procs (int): Number of gpus on single machine. |
|
launcher (str): How to launch distributed training. |
|
Only `slurm`, `pytorch` and `mpi` are available. |
|
|
|
Returns: |
|
List[Dict]: A list of tasks. |
|
""" |
|
|
|
tasks = [] |
|
for model, dataset, evaluator, load_from in zip( |
|
models, datasets, evaluators, load_froms): |
|
task = Config({ |
|
'model': model, |
|
'dataset': dataset, |
|
'evaluator': evaluator, |
|
'load_from': load_from, |
|
'work_dir': work_dir, |
|
'num_gpus': num_gpus, |
|
'num_procs': num_procs, |
|
'launcher': launcher |
|
}) |
|
tasks.append(task) |
|
|
|
return tasks |
|
|
|
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 list of dataset configs |
|
'evaluators': [], # a list of evaluator configs |
|
'load_froms': [], # a list of load_from paths |
|
'work_dir': '', # the work dir |
|
'num_gpus': int, # integer, number of gpus for each task |
|
'num_procs': int, # integer, number of gpus on single machine |
|
} |
|
|
|
Args: |
|
cfg (ConfigDict): The config dict, containing "models", "dataset" |
|
and "work_dir" keys. |
|
|
|
Returns: |
|
List[Dict]: A list of tasks. |
|
""" |
|
cfg = deepcopy(cfg) |
|
models = cfg['models'] |
|
datasets = cfg['datasets'] |
|
evaluators = cfg['evaluators'] |
|
load_froms = cfg['load_froms'] |
|
work_dir = cfg['work_dir'] |
|
num_gpus = cfg['num_gpus'] |
|
num_procs = cfg['num_procs'] |
|
launcher = cfg['launcher'] |
|
|
|
tasks = self.partition(models, datasets, evaluators, load_froms, |
|
work_dir, num_gpus, num_procs, launcher) |
|
|
|
self.logger.info(f'Partitioned into {len(tasks)} tasks.') |
|
for i, task in enumerate(tasks): |
|
model_name = task['model']['type'] |
|
dataset_name = task['dataset']['dataset']['type'] |
|
evaluator_name = task['evaluator'][0]['type'] |
|
self.logger.debug( |
|
f'Task {i}: {model_name}-{dataset_name}-{evaluator_name}') |
|
|
|
return tasks |
|
|
