init commit

.dev/clean_models.py

@@ -0,0 +1,125 @@
+import argparse
+import glob
+import json
+import os
+import os.path as osp
+
+import mmcv
+
+# build schedule look-up table to automatically find the final model
+SCHEDULES_LUT = {
+ '20ki': 20000,
+ '40ki': 40000,
+ '60ki': 60000,
+ '80ki': 80000,
+ '160ki': 160000
+}
+RESULTS_LUT = ['mIoU', 'mAcc', 'aAcc']
+
+
+def get_final_iter(config):
+ iter_num = SCHEDULES_LUT[config.split('_')[-2]]
+ return iter_num
+
+
+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')
+
+ args = parser.parse_args()
+ return args
+
+
+def main():
+ args = parse_args()
+ models_root = args.root
+ config_name = args.config
+
+ # 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)
+ 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
+ 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} not finished yet')
+ continue
+
+ # get logs
+ log_json_path = glob.glob(osp.join(exp_dir, '*.log.json'))[0]
+ log_txt_path = glob.glob(osp.join(exp_dir, '*.log'))[0]
+ model_performance = get_final_results(log_json_path, final_iter)
+
+ if model_performance is None:
+ print(f'{used_config} does not have performance')
+ continue
+
+ model_time = osp.split(log_txt_path)[-1].split('.')[0]
+ model_infos.append(
+ dict(
+ config=used_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
+ for model in model_infos:
+
+ model_name = osp.split(model['config'])[-1].split('.')[0]
+
+ model_name += '_' + model['model_time']
+ for checkpoints in mmcv.scandir(
+ osp.join(models_root, model['config']), suffix='.pth'):
+ if checkpoints.endswith(f"iter_{model['iters']}.pth"
+ ) or checkpoints.endswith('latest.pth'):
+ continue
+ print('removing {}'.format(
+ osp.join(models_root, model['config'], checkpoints)))
+ os.remove(osp.join(models_root, model['config'], checkpoints))
+
+
+if __name__ == '__main__':
+ main()

.dev/gather_models.py

@@ -0,0 +1,197 @@
+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()