gather_models.py

# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import glob
import hashlib
import json
import os
import os.path as osp
import shutil

import torch
from mmengine import Config
from mmengine.fileio import dump
from mmengine.utils import mkdir_or_exist, scandir

# build schedule look-up table to automatically find the final model
RESULTS_LUT = ['mIoU', 'mAcc', 'aAcc']


def calculate_file_sha256(file_path):
    """calculate file sha256 hash code."""
    with open(file_path, 'rb') as fp:
        sha256_cal = hashlib.sha256()
        sha256_cal.update(fp.read())
        return sha256_cal.hexdigest()


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)
    # The hash code calculation and rename command differ on different system
    # platform.
    sha = calculate_file_sha256(out_file)
    final_file = out_file.rstrip('.pth') + f'-{sha[:8]}.pth'
    os.rename(out_file, final_file)

    # Remove prefix and suffix
    final_file_name = osp.split(final_file)[1]
    final_file_name = osp.splitext(final_file_name)[0]

    return final_file_name


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()
    last_iter = 0
    with open(log_json_path) as f:
        for line in f.readlines():
            log_line = json.loads(line)
            if 'mode' not in log_line.keys():
                continue

            # When evaluation, the 'iter' of new log json is the evaluation
            # steps on single gpu.
            flag1 = ('aAcc' in log_line) or (log_line['mode'] == 'val')
            flag2 = (last_iter == iter_num - 50) or (last_iter == iter_num)
            if flag1 and flag2:
                result_dict.update({
                    key: log_line[key]
                    for key in RESULTS_LUT if key in log_line
                })
                return result_dict

            last_iter = log_line['iter']


def parse_args():
    parser = argparse.ArgumentParser(description='Gather benchmarked models')
    parser.add_argument(
        '-f', '--config-name', type=str, help='Process the selected config.')
    parser.add_argument(
        '-w',
        '--work-dir',
        default='work_dirs/',
        type=str,
        help='Ckpt storage root folder of benchmarked models to be gathered.')
    parser.add_argument(
        '-c',
        '--collect-dir',
        default='work_dirs/gather',
        type=str,
        help='Ckpt collect root folder of gathered models.')
    parser.add_argument(
        '--all', action='store_true', help='whether include .py and .log')

    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    work_dir = args.work_dir
    collect_dir = args.collect_dir
    selected_config_name = args.config_name
    mkdir_or_exist(collect_dir)

    # 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:
        config_name = osp.splitext(osp.basename(raw_config))[0]
        if osp.exists(osp.join(work_dir, config_name)):
            if (selected_config_name is None
                    or selected_config_name == config_name):
                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:
        config_name = osp.splitext(osp.basename(used_config))[0]
        exp_dir = osp.join(work_dir, config_name)
        # check whether the exps is finished
        final_iter = get_final_iter(used_config)
        final_model = f'iter_{final_iter}.pth'
        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_name=config_name,
                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:
        config_name = model['config_name']
        model_publish_dir = osp.join(collect_dir, config_name)

        publish_model_path = osp.join(model_publish_dir,
                                      config_name + '_' + model['model_time'])
        trained_model_path = osp.join(work_dir, config_name,
                                      '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:
            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'{config_name}_{model["log_json_path"]}'
        # copy log
        shutil.copy(
            osp.join(work_dir, config_name, 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(work_dir, config_name,
                         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('./configs', f'{config_name}.py')
            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)
    dump(models, osp.join(collect_dir, 'model_infos.json'), indent=4)


if __name__ == '__main__':
    main()