|
|
import argparse |
|
|
import copy |
|
|
import json |
|
|
import os |
|
|
|
|
|
import mmengine |
|
|
from mmengine.config import Config, ConfigDict |
|
|
|
|
|
from opencompass.utils import build_dataset_from_cfg, get_infer_output_path |
|
|
|
|
|
|
|
|
def parse_args(): |
|
|
parser = argparse.ArgumentParser( |
|
|
description='Merge patitioned predictions') |
|
|
parser.add_argument('config', help='Train config file path') |
|
|
parser.add_argument('-w', '--work-dir', default=None, type=str) |
|
|
parser.add_argument('-r', '--reuse', default='latest', type=str) |
|
|
parser.add_argument('-c', '--clean', action='store_true') |
|
|
parser.add_argument('-f', '--force', action='store_true') |
|
|
args = parser.parse_args() |
|
|
return args |
|
|
|
|
|
|
|
|
class PredictionMerger: |
|
|
|
|
|
def __init__(self, cfg: ConfigDict) -> None: |
|
|
self.cfg = cfg |
|
|
self.model_cfg = copy.deepcopy(self.cfg['model']) |
|
|
self.dataset_cfg = copy.deepcopy(self.cfg['dataset']) |
|
|
self.work_dir = self.cfg.get('work_dir') |
|
|
|
|
|
def run(self): |
|
|
filename = get_infer_output_path( |
|
|
self.model_cfg, self.dataset_cfg, |
|
|
os.path.join(self.work_dir, 'predictions')) |
|
|
root, ext = os.path.splitext(filename) |
|
|
partial_filename = root + '_0' + ext |
|
|
|
|
|
if os.path.exists( |
|
|
os.path.realpath(filename)) and not self.cfg['force']: |
|
|
return |
|
|
|
|
|
if not os.path.exists(os.path.realpath(partial_filename)): |
|
|
print(f'{filename} not found') |
|
|
return |
|
|
|
|
|
|
|
|
partial_filenames = [] |
|
|
preds, offset = {}, 0 |
|
|
i = 1 |
|
|
while os.path.exists(os.path.realpath(partial_filename)): |
|
|
partial_filenames.append(os.path.realpath(partial_filename)) |
|
|
_preds = mmengine.load(partial_filename) |
|
|
partial_filename = root + f'_{i}' + ext |
|
|
i += 1 |
|
|
for _o in range(len(_preds)): |
|
|
preds[str(offset)] = _preds[str(_o)] |
|
|
offset += 1 |
|
|
|
|
|
dataset = build_dataset_from_cfg(self.dataset_cfg) |
|
|
if len(preds) != len(dataset.test): |
|
|
print('length mismatch') |
|
|
return |
|
|
|
|
|
print(f'Merge {partial_filenames} to {filename}') |
|
|
with open(filename, 'w', encoding='utf-8') as f: |
|
|
json.dump(preds, f, indent=4, ensure_ascii=False) |
|
|
|
|
|
if self.cfg['clean']: |
|
|
for partial_filename in partial_filenames: |
|
|
print(f'Remove {partial_filename}') |
|
|
os.remove(partial_filename) |
|
|
|
|
|
|
|
|
def dispatch_tasks(cfg): |
|
|
for model in cfg['models']: |
|
|
for dataset in cfg['datasets']: |
|
|
PredictionMerger({ |
|
|
'model': model, |
|
|
'dataset': dataset, |
|
|
'work_dir': cfg['work_dir'], |
|
|
'clean': cfg['clean'], |
|
|
'force': cfg['force'], |
|
|
}).run() |
|
|
|
|
|
|
|
|
def main(): |
|
|
args = parse_args() |
|
|
cfg = Config.fromfile(args.config) |
|
|
|
|
|
if args.work_dir is not None: |
|
|
cfg['work_dir'] = args.work_dir |
|
|
else: |
|
|
cfg.setdefault('work_dir', './outputs/default') |
|
|
|
|
|
if args.reuse: |
|
|
if args.reuse == 'latest': |
|
|
if not os.path.exists(cfg.work_dir) or not os.listdir( |
|
|
cfg.work_dir): |
|
|
print('No previous results to reuse!') |
|
|
return |
|
|
else: |
|
|
dirs = os.listdir(cfg.work_dir) |
|
|
dir_time_str = sorted(dirs)[-1] |
|
|
else: |
|
|
dir_time_str = args.reuse |
|
|
cfg['work_dir'] = os.path.join(cfg.work_dir, dir_time_str) |
|
|
|
|
|
cfg['clean'] = args.clean |
|
|
cfg['force'] = args.force |
|
|
|
|
|
dispatch_tasks(cfg) |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
main() |
|
|
|
