# Copyright (c) OpenMMLab. All rights reserved. import argparse import glob import json import os.path as osp import shutil import subprocess from collections import OrderedDict import torch import yaml from mmengine.config import Config from mmengine.fileio import dump from mmengine.utils import mkdir_or_exist, scandir def ordered_yaml_dump(data, stream=None, Dumper=yaml.SafeDumper, **kwds): class OrderedDumper(Dumper): pass def _dict_representer(dumper, data): return dumper.represent_mapping( yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, data.items()) OrderedDumper.add_representer(OrderedDict, _dict_representer) return yaml.dump(data, stream, OrderedDumper, **kwds) def process_checkpoint(in_file, out_file): checkpoint = torch.load(in_file, map_location='cpu') # remove optimizer for smaller file size if 'optimizer' in checkpoint: del checkpoint['optimizer'] # remove ema state_dict for key in list(checkpoint['state_dict']): if key.startswith('ema_'): checkpoint['state_dict'].pop(key) # if it is necessary to remove some sensitive data in checkpoint['meta'], # add the code here. if torch.__version__ >= '1.6': torch.save(checkpoint, out_file, _use_new_zipfile_serialization=False) else: torch.save(checkpoint, out_file) sha = subprocess.check_output(['sha256sum', out_file]).decode() final_file = out_file.rstrip('.pth') + '-{}.pth'.format(sha[:8]) subprocess.Popen(['mv', out_file, final_file]) return final_file def is_by_epoch(config): cfg = Config.fromfile('./configs/' + config) return cfg.runner.type == 'EpochBasedRunner' def get_final_epoch_or_iter(config): cfg = Config.fromfile('./configs/' + config) if cfg.runner.type == 'EpochBasedRunner': return cfg.runner.max_epochs else: return cfg.runner.max_iters def get_best_epoch_or_iter(exp_dir): best_epoch_iter_full_path = list( sorted(glob.glob(osp.join(exp_dir, 'best_*.pth'))))[-1] best_epoch_or_iter_model_path = best_epoch_iter_full_path.split('/')[-1] best_epoch_or_iter = best_epoch_or_iter_model_path.\ split('_')[-1].split('.')[0] return best_epoch_or_iter_model_path, int(best_epoch_or_iter) def get_real_epoch_or_iter(config): cfg = Config.fromfile('./configs/' + config) if cfg.runner.type == 'EpochBasedRunner': epoch = cfg.runner.max_epochs if cfg.data.train.type == 'RepeatDataset': epoch *= cfg.data.train.times return epoch else: return cfg.runner.max_iters def get_final_results(log_json_path, epoch_or_iter, results_lut, by_epoch=True): result_dict = dict() last_val_line = None last_train_line = None last_val_line_idx = -1 last_train_line_idx = -1 with open(log_json_path, 'r') as f: for i, line in enumerate(f.readlines()): log_line = json.loads(line) if 'mode' not in log_line.keys(): continue if by_epoch: if (log_line['mode'] == 'train' and log_line['epoch'] == epoch_or_iter): result_dict['memory'] = log_line['memory'] if (log_line['mode'] == 'val' and log_line['epoch'] == epoch_or_iter): result_dict.update({ key: log_line[key] for key in results_lut if key in log_line }) return result_dict else: if log_line['mode'] == 'train': last_train_line_idx = i last_train_line = log_line if log_line and log_line['mode'] == 'val': last_val_line_idx = i last_val_line = log_line # bug: max_iters = 768, last_train_line['iter'] = 750 assert last_val_line_idx == last_train_line_idx + 1, \ 'Log file is incomplete' result_dict['memory'] = last_train_line['memory'] result_dict.update({ key: last_val_line[key] for key in results_lut if key in last_val_line }) return result_dict def get_dataset_name(config): # If there are more dataset, add here. name_map = dict( CityscapesDataset='Cityscapes', CocoDataset='COCO', CocoPanopticDataset='COCO', DeepFashionDataset='Deep Fashion', LVISV05Dataset='LVIS v0.5', LVISV1Dataset='LVIS v1', VOCDataset='Pascal VOC', WIDERFaceDataset='WIDER Face', OpenImagesDataset='OpenImagesDataset', OpenImagesChallengeDataset='OpenImagesChallengeDataset', Objects365V1Dataset='Objects365 v1', Objects365V2Dataset='Objects365 v2') cfg = Config.fromfile('./configs/' + config) return name_map[cfg.dataset_type] def convert_model_info_to_pwc(model_infos): pwc_files = {} for model in model_infos: cfg_folder_name = osp.split(model['config'])[-2] pwc_model_info = OrderedDict() pwc_model_info['Name'] = osp.split(model['config'])[-1].split('.')[0] pwc_model_info['In Collection'] = 'Please fill in Collection name' pwc_model_info['Config'] = osp.join('configs', model['config']) # get metadata memory = round(model['results']['memory'] / 1024, 1) meta_data = OrderedDict() meta_data['Training Memory (GB)'] = memory if 'epochs' in model: meta_data['Epochs'] = get_real_epoch_or_iter(model['config']) else: meta_data['Iterations'] = get_real_epoch_or_iter(model['config']) pwc_model_info['Metadata'] = meta_data # get dataset name dataset_name = get_dataset_name(model['config']) # get results results = [] # if there are more metrics, add here. if 'bbox_mAP' in model['results']: metric = round(model['results']['bbox_mAP'] * 100, 1) results.append( OrderedDict( Task='Object Detection', Dataset=dataset_name, Metrics={'box AP': metric})) if 'segm_mAP' in model['results']: metric = round(model['results']['segm_mAP'] * 100, 1) results.append( OrderedDict( Task='Instance Segmentation', Dataset=dataset_name, Metrics={'mask AP': metric})) if 'PQ' in model['results']: metric = round(model['results']['PQ'], 1) results.append( OrderedDict( Task='Panoptic Segmentation', Dataset=dataset_name, Metrics={'PQ': metric})) pwc_model_info['Results'] = results link_string = 'https://download.openmmlab.com/mmdetection/v2.0/' link_string += '{}/{}'.format(model['config'].rstrip('.py'), osp.split(model['model_path'])[-1]) pwc_model_info['Weights'] = link_string if cfg_folder_name in pwc_files: pwc_files[cfg_folder_name].append(pwc_model_info) else: pwc_files[cfg_folder_name] = [pwc_model_info] return pwc_files def parse_args(): parser = argparse.ArgumentParser(description='Gather benchmarked models') parser.add_argument( 'root', type=str, help='root path of benchmarked models to be gathered') parser.add_argument( 'out', type=str, help='output path of gathered models to be stored') parser.add_argument( '--best', action='store_true', help='whether to gather the best model.') args = parser.parse_args() return args def main(): args = parse_args() models_root = args.root models_out = args.out mkdir_or_exist(models_out) # find all models in the root directory to be gathered raw_configs = list(scandir('./configs', '.py', recursive=True)) # filter configs that is not trained in the experiments dir used_configs = [] for raw_config in raw_configs: if osp.exists(osp.join(models_root, raw_config)): used_configs.append(raw_config) print(f'Find {len(used_configs)} models to be gathered') # find final_ckpt and log file for trained each config # and parse the best performance model_infos = [] for used_config in used_configs: exp_dir = osp.join(models_root, used_config) by_epoch = is_by_epoch(used_config) # check whether the exps is finished if args.best is True: final_model, final_epoch_or_iter = get_best_epoch_or_iter(exp_dir) else: final_epoch_or_iter = get_final_epoch_or_iter(used_config) final_model = '{}_{}.pth'.format('epoch' if by_epoch else 'iter', final_epoch_or_iter) model_path = osp.join(exp_dir, final_model) # skip if the model is still training if not osp.exists(model_path): continue # get the latest logs log_json_path = list( sorted(glob.glob(osp.join(exp_dir, '*.log.json'))))[-1] log_txt_path = list(sorted(glob.glob(osp.join(exp_dir, '*.log'))))[-1] cfg = Config.fromfile('./configs/' + used_config) results_lut = cfg.evaluation.metric if not isinstance(results_lut, list): results_lut = [results_lut] # case when using VOC, the evaluation key is only 'mAP' # when using Panoptic Dataset, the evaluation key is 'PQ'. for i, key in enumerate(results_lut): if 'mAP' not in key and 'PQ' not in key: results_lut[i] = key + '_mAP' model_performance = get_final_results(log_json_path, final_epoch_or_iter, results_lut, by_epoch) if model_performance is None: continue model_time = osp.split(log_txt_path)[-1].split('.')[0] model_info = dict( config=used_config, results=model_performance, model_time=model_time, final_model=final_model, log_json_path=osp.split(log_json_path)[-1]) model_info['epochs' if by_epoch else 'iterations'] =\ final_epoch_or_iter model_infos.append(model_info) # publish model for each checkpoint publish_model_infos = [] for model in model_infos: model_publish_dir = osp.join(models_out, model['config'].rstrip('.py')) mkdir_or_exist(model_publish_dir) model_name = osp.split(model['config'])[-1].split('.')[0] model_name += '_' + model['model_time'] publish_model_path = osp.join(model_publish_dir, model_name) trained_model_path = osp.join(models_root, model['config'], model['final_model']) # convert model final_model_path = process_checkpoint(trained_model_path, publish_model_path) # copy log shutil.copy( osp.join(models_root, model['config'], model['log_json_path']), osp.join(model_publish_dir, f'{model_name}.log.json')) shutil.copy( osp.join(models_root, model['config'], model['log_json_path'].rstrip('.json')), osp.join(model_publish_dir, f'{model_name}.log')) # copy config to guarantee reproducibility config_path = model['config'] config_path = osp.join( 'configs', config_path) if 'configs' not in config_path else config_path target_config_path = osp.split(config_path)[-1] shutil.copy(config_path, osp.join(model_publish_dir, target_config_path)) model['model_path'] = final_model_path publish_model_infos.append(model) models = dict(models=publish_model_infos) print(f'Totally gathered {len(publish_model_infos)} models') dump(models, osp.join(models_out, 'model_info.json')) pwc_files = convert_model_info_to_pwc(publish_model_infos) for name in pwc_files: with open(osp.join(models_out, name + '_metafile.yml'), 'w') as f: ordered_yaml_dump(pwc_files[name], f, encoding='utf-8') if __name__ == '__main__': main()