Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support to collect the best models #6560

Merged
merged 6 commits into from
Nov 24, 2021
Merged

Support to collect the best models #6560

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c5e89d1
Fix mosaic repr typo (#6523)
lkm2835 Nov 18, 2021
a3258bc
Include mmflow in readme (#6545)
Czm369 Nov 19, 2021
25c8d75
Make OHEM work with seesaw loss (#6514)
ohwi Nov 19, 2021
6e7a496
update
hhaAndroid Nov 22, 2021
2d80350
support gather best model
hhaAndroid Nov 22, 2021
7645922
Merge branch 'dev-v2.19.0' of https://github.com/open-mmlab/mmdetecti…
hhaAndroid Nov 22, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 23 additions & 7 deletions .dev_scripts/gather_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,14 @@ def get_final_epoch(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
Expand Down Expand Up @@ -160,6 +168,10 @@ def parse_args():
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
Expand Down Expand Up @@ -187,10 +199,13 @@ def main():
for used_config in used_configs:
exp_dir = osp.join(models_root, used_config)
# check whether the exps is finished
final_epoch = get_final_epoch(used_config)
final_model = 'epoch_{}.pth'.format(final_epoch)
model_path = osp.join(exp_dir, final_model)
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
Expand Down Expand Up @@ -221,6 +236,7 @@ def main():
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
Expand All @@ -234,7 +250,7 @@ def main():
model_name += '_' + model['model_time']
publish_model_path = osp.join(model_publish_dir, model_name)
trained_model_path = osp.join(models_root, model['config'],
'epoch_{}.pth'.format(model['epochs']))
model['final_model'])

# convert model
final_model_path = process_checkpoint(trained_model_path,
Expand All @@ -254,9 +270,9 @@ def main():
config_path = osp.join(
'configs',
config_path) if 'configs' not in config_path else config_path
target_cconfig_path = osp.split(config_path)[-1]
shutil.copy(config_path,
osp.join(model_publish_dir, target_cconfig_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)
Expand Down