198 lines
6.8 KiB
Python
198 lines
6.8 KiB
Python
|
import argparse
|
||
|
|
import glob
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import os.path as osp
|
||
|
|
import shutil
|
||
|
|
import subprocess
|
||
|
|
|
||
|
|
import mmcv
|
||
|
|
import torch
|
||
|
|
|
||
|
|
# build schedule look-up table to automatically find the final model
|
||
|
|
RESULTS_LUT = ['mIoU', 'mAcc', 'aAcc']
|
||
|
|
|
||
|
|
|
||
|
|
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']
|
||
|
|
# if it is necessary to remove some sensitive data in checkpoint['meta'],
|
||
|
|
# add the code here.
|
||
|
|
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_iter(config):
|
||
|
|
iter_num = config.split('_')[-2]
|
||
|
|
assert iter_num.endswith('k')
|
||
|
|
return int(iter_num[:-1]) * 1000
|
||
|
|
|
||
|
|
|
||
|
|
def get_final_results(log_json_path, iter_num):
|
||
|
|
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['iter'] == iter_num:
|
||
|
|
result_dict['memory'] = log_line['memory']
|
||
|
|
|
||
|
|
if log_line['iter'] == iter_num:
|
||
|
|
result_dict.update({
|
||
|
|
key: log_line[key]
|
||
|
|
for key in RESULTS_LUT if key in log_line
|
||
|
|
})
|
||
|
|
return result_dict
|
||
|
|
|
||
|
|
|
||
|
|
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(
|
||
|
|
'config',
|
||
|
|
type=str,
|
||
|
|
help='root path of benchmarked configs to be gathered')
|
||
|
|
parser.add_argument(
|
||
|
|
'out_dir',
|
||
|
|
type=str,
|
||
|
|
help='output path of gathered models to be stored')
|
||
|
|
parser.add_argument('out_file', type=str, help='the output json file name')
|
||
|
|
parser.add_argument(
|
||
|
|
'--filter', type=str, nargs='+', default=[], help='config filter')
|
||
|
|
parser.add_argument(
|
||
|
|
'--all', action='store_true', help='whether include .py and .log')
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
return args
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
args = parse_args()
|
||
|
|
models_root = args.root
|
||
|
|
models_out = args.out_dir
|
||
|
|
config_name = args.config
|
||
|
|
mmcv.mkdir_or_exist(models_out)
|
||
|
|
|
||
|
|
# find all models in the root directory to be gathered
|
||
|
|
raw_configs = list(mmcv.scandir(config_name, '.py', recursive=True))
|
||
|
|
|
||
|
|
# filter configs that is not trained in the experiments dir
|
||
|
|
used_configs = []
|
||
|
|
for raw_config in raw_configs:
|
||
|
|
work_dir = osp.splitext(osp.basename(raw_config))[0]
|
||
|
|
if osp.exists(osp.join(models_root, work_dir)):
|
||
|
|
used_configs.append((work_dir, 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, raw_config in used_configs:
|
||
|
|
bypass = True
|
||
|
|
for p in args.filter:
|
||
|
|
if p in used_config:
|
||
|
|
bypass = False
|
||
|
|
break
|
||
|
|
if bypass:
|
||
|
|
continue
|
||
|
|
exp_dir = osp.join(models_root, used_config)
|
||
|
|
# check whether the exps is finished
|
||
|
|
final_iter = get_final_iter(used_config)
|
||
|
|
final_model = 'iter_{}.pth'.format(final_iter)
|
||
|
|
model_path = osp.join(exp_dir, final_model)
|
||
|
|
|
||
|
|
# skip if the model is still training
|
||
|
|
if not osp.exists(model_path):
|
||
|
|
print(f'{used_config} train not finished yet')
|
||
|
|
continue
|
||
|
|
|
||
|
|
# get logs
|
||
|
|
log_json_paths = glob.glob(osp.join(exp_dir, '*.log.json'))
|
||
|
|
log_json_path = log_json_paths[0]
|
||
|
|
model_performance = None
|
||
|
|
for idx, _log_json_path in enumerate(log_json_paths):
|
||
|
|
model_performance = get_final_results(_log_json_path, final_iter)
|
||
|
|
if model_performance is not None:
|
||
|
|
log_json_path = _log_json_path
|
||
|
|
break
|
||
|
|
|
||
|
|
if model_performance is None:
|
||
|
|
print(f'{used_config} model_performance is None')
|
||
|
|
continue
|
||
|
|
|
||
|
|
model_time = osp.split(log_json_path)[-1].split('.')[0]
|
||
|
|
model_infos.append(
|
||
|
|
dict(
|
||
|
|
config=used_config,
|
||
|
|
raw_config=raw_config,
|
||
|
|
results=model_performance,
|
||
|
|
iters=final_iter,
|
||
|
|
model_time=model_time,
|
||
|
|
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['raw_config'].rstrip('.py'))
|
||
|
|
model_name = osp.split(model['config'])[-1].split('.')[0]
|
||
|
|
|
||
|
|
publish_model_path = osp.join(model_publish_dir,
|
||
|
|
model_name + '_' + model['model_time'])
|
||
|
|
trained_model_path = osp.join(models_root, model['config'],
|
||
|
|
'iter_{}.pth'.format(model['iters']))
|
||
|
|
if osp.exists(model_publish_dir):
|
||
|
|
for file in os.listdir(model_publish_dir):
|
||
|
|
if file.endswith('.pth'):
|
||
|
|
print(f'model {file} found')
|
||
|
|
model['model_path'] = osp.abspath(
|
||
|
|
osp.join(model_publish_dir, file))
|
||
|
|
break
|
||
|
|
if 'model_path' not in model:
|
||
|
|
print(f'dir {model_publish_dir} exists, no model found')
|
||
|
|
|
||
|
|
else:
|
||
|
|
mmcv.mkdir_or_exist(model_publish_dir)
|
||
|
|
|
||
|
|
# convert model
|
||
|
|
final_model_path = process_checkpoint(trained_model_path,
|
||
|
|
publish_model_path)
|
||
|
|
model['model_path'] = final_model_path
|
||
|
|
|
||
|
|
new_json_path = f'{model_name}-{model["log_json_path"]}'
|
||
|
|
# copy log
|
||
|
|
shutil.copy(
|
||
|
|
osp.join(models_root, model['config'], model['log_json_path']),
|
||
|
|
osp.join(model_publish_dir, new_json_path))
|
||
|
|
if args.all:
|
||
|
|
new_txt_path = new_json_path.rstrip('.json')
|
||
|
|
shutil.copy(
|
||
|
|
osp.join(models_root, model['config'],
|
||
|
|
model['log_json_path'].rstrip('.json')),
|
||
|
|
osp.join(model_publish_dir, new_txt_path))
|
||
|
|
|
||
|
|
if args.all:
|
||
|
|
# copy config to guarantee reproducibility
|
||
|
|
raw_config = osp.join(config_name, model['raw_config'])
|
||
|
|
mmcv.Config.fromfile(raw_config).dump(
|
||
|
|
osp.join(model_publish_dir, osp.basename(raw_config)))
|
||
|
|
|
||
|
|
publish_model_infos.append(model)
|
||
|
|
|
||
|
|
models = dict(models=publish_model_infos)
|
||
|
|
mmcv.dump(models, osp.join(models_out, args.out_file))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|