# Copyright (c) OpenMMLab. All rights reserved. import argparse import os import warnings from numbers import Number import mmcv import numpy as np import torch from mmcv import DictAction from mmcv.parallel import MMDataParallel, MMDistributedDataParallel from mmcv.runner import get_dist_info, init_dist, load_checkpoint from mmcls.apis import multi_gpu_test, single_gpu_test from mmcls.datasets import build_dataloader, build_dataset from mmcls.models import build_classifier # TODO import `wrap_fp16_model` from mmcv and delete them from mmcls try: from mmcv.runner import wrap_fp16_model except ImportError: warnings.warn('wrap_fp16_model from mmcls will be deprecated.' 'Please install mmcv>=1.1.4.') from mmcls.core import wrap_fp16_model def parse_args(): parser = argparse.ArgumentParser(description='mmcls test model') parser.add_argument('config', help='test config file path') parser.add_argument('checkpoint', help='checkpoint file') parser.add_argument('--out', help='output result file') out_options = ['class_scores', 'pred_score', 'pred_label', 'pred_class'] parser.add_argument( '--out-items', nargs='+', default=['all'], choices=out_options + ['none', 'all'], help='Besides metrics, what items will be included in the output ' f'result file. You can choose some of ({", ".join(out_options)}), ' 'or use "all" to include all above, or use "none" to disable all of ' 'above. Defaults to output all.', metavar='') parser.add_argument( '--metrics', type=str, nargs='+', help='evaluation metrics, which depends on the dataset, e.g., ' '"accuracy", "precision", "recall", "f1_score", "support" for single ' 'label dataset, and "mAP", "CP", "CR", "CF1", "OP", "OR", "OF1" for ' 'multi-label dataset') parser.add_argument('--show', action='store_true', help='show results') parser.add_argument( '--show-dir', help='directory where painted images will be saved') parser.add_argument( '--gpu-collect', action='store_true', help='whether to use gpu to collect results') parser.add_argument('--tmpdir', help='tmp dir for writing some results') parser.add_argument( '--cfg-options', nargs='+', action=DictAction, help='override some settings in the used config, the key-value pair ' 'in xxx=yyy format will be merged into config file. If the value to ' 'be overwritten is a list, it should be like key="[a,b]" or key=a,b ' 'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" ' 'Note that the quotation marks are necessary and that no white space ' 'is allowed.') parser.add_argument( '--options', nargs='+', action=DictAction, help='override some settings in the used config, the key-value pair ' 'in xxx=yyy format will be merged into config file (deprecate), ' 'change to --cfg-options instead.') parser.add_argument( '--metric-options', nargs='+', action=DictAction, default={}, help='custom options for evaluation, the key-value pair in xxx=yyy ' 'format will be parsed as a dict metric_options for dataset.evaluate()' ' function.') parser.add_argument( '--show-options', nargs='+', action=DictAction, help='custom options for show_result. key-value pair in xxx=yyy.' 'Check available options in `model.show_result`.') parser.add_argument( '--launcher', choices=['none', 'pytorch', 'slurm', 'mpi'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) parser.add_argument( '--device', choices=['cpu', 'cuda'], default='cuda', help='device used for testing') args = parser.parse_args() if 'LOCAL_RANK' not in os.environ: os.environ['LOCAL_RANK'] = str(args.local_rank) if args.options and args.cfg_options: raise ValueError( '--options and --cfg-options cannot be both ' 'specified, --options is deprecated in favor of --cfg-options') if args.options: warnings.warn('--options is deprecated in favor of --cfg-options') args.cfg_options = args.options return args def main(): args = parse_args() cfg = mmcv.Config.fromfile(args.config) if args.cfg_options is not None: cfg.merge_from_dict(args.cfg_options) # set cudnn_benchmark if cfg.get('cudnn_benchmark', False): torch.backends.cudnn.benchmark = True cfg.model.pretrained = None cfg.data.test.test_mode = True assert args.metrics or args.out, \ 'Please specify at least one of output path and evaluation metrics.' # init distributed env first, since logger depends on the dist info. if args.launcher == 'none': distributed = False else: distributed = True init_dist(args.launcher, **cfg.dist_params) # build the dataloader dataset = build_dataset(cfg.data.test) # the extra round_up data will be removed during gpu/cpu collect data_loader = build_dataloader( dataset, samples_per_gpu=cfg.data.samples_per_gpu, workers_per_gpu=cfg.data.workers_per_gpu, dist=distributed, shuffle=False, round_up=True) # build the model and load checkpoint model = build_classifier(cfg.model) fp16_cfg = cfg.get('fp16', None) if fp16_cfg is not None: wrap_fp16_model(model) checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu') if 'CLASSES' in checkpoint.get('meta', {}): CLASSES = checkpoint['meta']['CLASSES'] else: from mmcls.datasets import ImageNet warnings.simplefilter('once') warnings.warn('Class names are not saved in the checkpoint\'s ' 'meta data, use imagenet by default.') CLASSES = ImageNet.CLASSES if not distributed: if args.device == 'cpu': model = model.cpu() else: model = MMDataParallel(model, device_ids=[0]) model.CLASSES = CLASSES show_kwargs = {} if args.show_options is None else args.show_options outputs = single_gpu_test(model, data_loader, args.show, args.show_dir, **show_kwargs) else: model = MMDistributedDataParallel( model.cuda(), device_ids=[torch.cuda.current_device()], broadcast_buffers=False) outputs = multi_gpu_test(model, data_loader, args.tmpdir, args.gpu_collect) rank, _ = get_dist_info() if rank == 0: results = {} if args.metrics: eval_results = dataset.evaluate(outputs, args.metrics, args.metric_options) results.update(eval_results) for k, v in eval_results.items(): if isinstance(v, np.ndarray): v = [round(out, 2) for out in v.tolist()] elif isinstance(v, Number): v = round(v, 2) else: raise ValueError(f'Unsupport metric type: {type(v)}') print(f'\n{k} : {v}') if args.out: if 'none' not in args.out_items: scores = np.vstack(outputs) pred_score = np.max(scores, axis=1) pred_label = np.argmax(scores, axis=1) pred_class = [CLASSES[lb] for lb in pred_label] res_items = { 'class_scores': scores, 'pred_score': pred_score, 'pred_label': pred_label, 'pred_class': pred_class } if 'all' in args.out_items: results.update(res_items) else: for key in args.out_items: results[key] = res_items[key] print(f'\ndumping results to {args.out}') mmcv.dump(results, args.out) if __name__ == '__main__': main()