292 lines
10 KiB
Python
292 lines
10 KiB
Python
# 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 mmcv
|
|
import torch
|
|
import yaml
|
|
|
|
|
|
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 get_final_epoch(config):
|
|
cfg = mmcv.Config.fromfile('./configs/' + config)
|
|
return cfg.runner.max_epochs
|
|
|
|
|
|
def get_best_epoch(exp_dir):
|
|
best_epoch_full_path = list(
|
|
sorted(glob.glob(osp.join(exp_dir, 'best_*.pth'))))[-1]
|
|
best_epoch_model_path = best_epoch_full_path.split('/')[-1]
|
|
best_epoch = best_epoch_model_path.split('_')[-1].split('.')[0]
|
|
return best_epoch_model_path, int(best_epoch)
|
|
|
|
|
|
def get_real_epoch(config):
|
|
cfg = mmcv.Config.fromfile('./configs/' + config)
|
|
epoch = cfg.runner.max_epochs
|
|
if cfg.data.train.type == 'RepeatDataset':
|
|
epoch *= cfg.data.train.times
|
|
return epoch
|
|
|
|
|
|
def get_final_results(log_json_path, epoch, results_lut):
|
|
result_dict = dict()
|
|
with open(log_json_path, 'r') as f:
|
|
for line in f.readlines():
|
|
log_line = json.loads(line)
|
|
if 'mode' not in log_line.keys():
|
|
continue
|
|
|
|
if log_line['mode'] == 'train' and log_line['epoch'] == epoch:
|
|
result_dict['memory'] = log_line['memory']
|
|
|
|
if log_line['mode'] == 'val' and log_line['epoch'] == epoch:
|
|
result_dict.update({
|
|
key: log_line[key]
|
|
for key in results_lut if key in log_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')
|
|
cfg = mmcv.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)
|
|
epochs = get_real_epoch(model['config'])
|
|
meta_data = OrderedDict()
|
|
meta_data['Training Memory (GB)'] = memory
|
|
meta_data['Epochs'] = epochs
|
|
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
|
|
mmcv.mkdir_or_exist(models_out)
|
|
|
|
# find all models in the root directory to be gathered
|
|
raw_configs = list(mmcv.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)
|
|
# check whether the exps is finished
|
|
if args.best is True:
|
|
final_model, final_epoch = get_best_epoch(exp_dir)
|
|
else:
|
|
final_epoch = get_final_epoch(used_config)
|
|
final_model = 'epoch_{}.pth'.format(final_epoch)
|
|
|
|
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 = mmcv.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 + 'm_AP'
|
|
model_performance = get_final_results(log_json_path, final_epoch,
|
|
results_lut)
|
|
|
|
if model_performance is None:
|
|
continue
|
|
|
|
model_time = osp.split(log_txt_path)[-1].split('.')[0]
|
|
model_infos.append(
|
|
dict(
|
|
config=used_config,
|
|
results=model_performance,
|
|
epochs=final_epoch,
|
|
model_time=model_time,
|
|
final_model=final_model,
|
|
log_json_path=osp.split(log_json_path)[-1]))
|
|
|
|
# publish model for each checkpoint
|
|
publish_model_infos = []
|
|
for model in model_infos:
|
|
model_publish_dir = osp.join(models_out, model['config'].rstrip('.py'))
|
|
mmcv.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')
|
|
mmcv.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()
|