test.py

"""Testing entry for the GrokDet lightweight runtime.

Examples:
    python test.py CONFIG CHECKPOINT --device cuda:0
    python test.py CONFIG CHECKPOINT --out predictions.pkl
    python test.py CONFIG CHECKPOINT --show-dir predictions
    torchrun --nproc_per_node=2 test.py CONFIG CHECKPOINT --launcher pytorch

The checkpoint can be omitted for composite tracking configs whose components
define their own initialization checkpoints.
"""

import argparse
import logging
import os
import os.path as osp
from datetime import datetime

from datasets.samplers import TrackImgSampler
from types import SimpleNamespace

import torch
from mmengine.config import Config, DictAction
from mmengine.evaluator import BaseMetric, Evaluator
from mmengine.logging import MMLogger, MessageHub
from mmengine.runner import load_checkpoint as load_model_checkpoint

from runtime.builders import (build_dataloader_from_cfg, build_log_processor,
                              build_visualizer, process_eval_batch,
                              reset_evaluator_results, validate_metric_cfg)
from runtime.common import require_dataset_meta, to_plain_dict
from runtime.context import RuntimeContext
from runtime.dist import (barrier, init_distributed_runtime, resolve_device,
                          sync_string_from_rank0,
                          wrap_model_for_distributed_training)
from runtime.hooks import build_runtime_hooks
from train import (ensure_non_empty_dataloader, setup_randomness,
                   setup_runtime_environment, validate_dataloader_cfg)
from utils.builder import build_dataset, build_evaluator, build_model


def parse_args():
    parser = argparse.ArgumentParser(
        description='Test and evaluate a GrokDet model')
    parser.add_argument('config', help='test config file path')
    parser.add_argument(
        'checkpoint', nargs='?', default=None,
        help='model checkpoint; omit to use model/submodule init_cfg')
    parser.add_argument('--work-dir', help='directory for test artifacts')
    parser.add_argument('--out', help='dump predictions to a pickle file')
    parser.add_argument(
        '--show', action='store_true', help='display prediction results')
    parser.add_argument(
        '--show-dir', help='directory for rendered predictions')
    parser.add_argument(
        '--wait-time', type=float, default=2.0,
        help='display interval used with --show')
    parser.add_argument(
        '--device', default='cuda',
        help='single-process device; distributed testing uses LOCAL_RANK')
    parser.add_argument(
        '--launcher', choices=('none', 'pytorch', 'slurm', 'mpi'),
        default='none', help='distributed job launcher')
    parser.add_argument(
        '--cfg-options', nargs='+', action=DictAction,
        help='override config values with key=value pairs')
    parser.add_argument(
        '--strict', action='store_true',
        help='require exact checkpoint key matching')
    parser.add_argument(
        '--tta', action='store_true',
        help='request test-time augmentation (not supported yet)')
    parser.add_argument(
        '--local-rank', '--local_rank', dest='local_rank', type=int,
        default=0, help=argparse.SUPPRESS)
    args = parser.parse_args()
    os.environ.setdefault('LOCAL_RANK', str(args.local_rank))
    if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
        parser.error('--out must end with .pkl or .pickle')
    if args.tta:
        parser.error(
            '--tta is not supported by the GrokDet explicit runtime yet')
    return args


def prepare_test_work_dir(cfg, config_path, cli_work_dir):
    if cli_work_dir is not None:
        work_dir = cli_work_dir
    elif cfg.get('work_dir', None):
        work_dir = cfg.work_dir
    else:
        name = osp.splitext(osp.basename(config_path))[0]
        work_dir = osp.join('work_dirs', name)
    cfg.work_dir = osp.abspath(work_dir)
    os.makedirs(cfg.work_dir, exist_ok=True)


def build_test_logger(work_dir, timestamp):
    return MMLogger.get_instance(
        name='GrokDetTest',
        logger_name='GrokDetTest',
        log_file=osp.join(work_dir, f'test_{timestamp}.log'),
        log_level=logging.INFO,
        distributed=False)


def configure_visualization(cfg, args):
    if not (args.show or args.show_dir):
        return
    hooks = cfg.get('default_hooks', None)
    visualization = hooks.get('visualization') if isinstance(hooks, dict) else None
    if not isinstance(visualization, dict):
        raise KeyError(
            '--show/--show-dir requires default_hooks.visualization')
    visualization['draw'] = True
    visualization['show'] = args.show
    visualization['wait_time'] = args.wait_time
    if args.show_dir is not None:
        visualization['test_out_dir'] = args.show_dir


def build_test_data(cfg, logger):
    loader_cfg = to_plain_dict(cfg.get('test_dataloader', None))
    if loader_cfg is None:
        raise KeyError('Config must define test_dataloader.')
    validate_dataloader_cfg(loader_cfg, 'TestLoop', 'test_dataloader')
    dataset_cfg = loader_cfg.get('dataset', None)
    if not isinstance(dataset_cfg, dict):
        raise KeyError('test_dataloader.dataset must be a config dict.')
    logger.info('Building test dataset and dataloader...')
    dataset = build_dataset(dataset_cfg)
    dataloader = build_dataloader_from_cfg(
        loader_cfg, dataset, default_shuffle=False)
    ensure_non_empty_dataloader(dataloader, 'test_dataloader')
    return dataset, dataloader


def append_dump_metric(evaluator, out_file):
    if out_file is None:
        return evaluator
    out_file = osp.abspath(out_file)
    os.makedirs(osp.dirname(out_file), exist_ok=True)
    dump_metric = build_evaluator(
        dict(type='DumpDetResults', out_file_path=out_file))
    if evaluator is None:
        return dump_metric
    if isinstance(evaluator, BaseMetric):
        return Evaluator([evaluator, dump_metric])
    if isinstance(evaluator, Evaluator):
        evaluator.metrics.append(dump_metric)
        return evaluator
    raise TypeError(f'Unsupported evaluator type: {type(evaluator)}')


def build_test_evaluator(cfg, out_file):
    evaluator_cfg = to_plain_dict(cfg.get('test_evaluator', None))
    evaluator = None
    if evaluator_cfg is not None:
        validate_metric_cfg(evaluator_cfg)
        if isinstance(evaluator_cfg, list):
            metrics = []
            for metric_cfg in evaluator_cfg:
                metric = build_evaluator(metric_cfg)
                if isinstance(metric, Evaluator):
                    raise NotImplementedError(
                        'Nested test evaluators are not supported.')
                metrics.append(metric)
            evaluator = Evaluator(metrics)
        elif isinstance(evaluator_cfg, dict):
            evaluator = build_evaluator(evaluator_cfg.copy())
        else:
            raise TypeError(
                'test_evaluator must be a dict, list or None, but got '
                f'{type(evaluator_cfg)}')
    return append_dump_metric(evaluator, out_file)


def build_test_model(cfg, checkpoint, strict, logger, distributed):
    logger.info('Building model...')
    model = build_model(cfg.model).to(cfg.device)
    if checkpoint is None:
        logger.info(
            'No top-level checkpoint supplied; using model init_cfg.')
        if hasattr(model, 'init_weights'):
            model.init_weights()
        loaded_checkpoint = None
    else:
        checkpoint = (osp.abspath(checkpoint)
                      if osp.isfile(checkpoint) else checkpoint)
        logger.info(f'Loading checkpoint: {checkpoint}')
        loaded_checkpoint = load_model_checkpoint(
            model, checkpoint, map_location='cpu', strict=strict,
            logger=logger)
    if distributed:
        model = wrap_model_for_distributed_training(model, cfg)
    return model, loaded_checkpoint


def build_test_context(cfg, logger, model, dataset, dataloader, timestamp):
    dataset_meta = require_dataset_meta(
        dataset, dataloader, 'test_dataloader.dataset')
    visualizer = build_visualizer(
        cfg, work_dir=cfg.work_dir, dataset_meta=dataset_meta)
    hook_manager = build_runtime_hooks(cfg, logger)
    context = RuntimeContext(
        model=model, cfg=cfg, work_dir=cfg.work_dir, logger=logger,
        message_hub=MessageHub.get_instance('grokdet_test'),
        optim_wrapper=None, param_schedulers=[],
        log_processor=build_log_processor(cfg), visualizer=visualizer,
        train_dataloader=None, val_dataloader=None, val_dataset=None,
        hook_manager=hook_manager, loop_type='TestLoop', val_interval=1,
        val_metric_key=None, max_epochs=0, max_iters=0,
        timestamp=timestamp)
    context.test_dataloader = dataloader
    context.test_dataset = dataset
    context._train_loop = None
    context.train_loop = None
    context.test_loop = SimpleNamespace(dataloader=dataloader)
    context._test_loop = context.test_loop
    return context


def run_test(runner, evaluator):
    dataset_meta = require_dataset_meta(
        runner.test_dataset, runner.test_dataloader,
        'test_dataloader.dataset')
    if evaluator is not None and hasattr(evaluator, 'dataset_meta'):
        evaluator.dataset_meta = dataset_meta
        reset_evaluator_results(evaluator)

    runner.model.eval()
    runner.hook_manager.call('before_test', runner)
    runner.hook_manager.call('before_test_epoch', runner)
    try:
        with torch.inference_mode():
            for batch_idx, data_batch in enumerate(runner.test_dataloader):
                runner.iter = batch_idx
                runner.hook_manager.call(
                    'before_test_iter', runner, batch_idx=batch_idx,
                    data_batch=data_batch)
                outputs = runner.model.test_step(data_batch)
                if evaluator is not None:
                    process_eval_batch(evaluator, data_batch, outputs)
                runner.hook_manager.call(
                    'after_test_iter', runner, batch_idx=batch_idx,
                    data_batch=data_batch, outputs=outputs)
        metrics = {}
        if evaluator is not None:
            if isinstance(runner.test_dataloader.sampler, TrackImgSampler):
                evaluation_size = runner.test_dataset.num_all_imgs
            else:
                evaluation_size = len(runner.test_dataset)
            metrics = evaluator.evaluate(size=evaluation_size)
        for name, value in metrics.items():
            if isinstance(value, (int, float)):
                runner.message_hub.update_scalar(f'test/{name}', float(value))
        runner.hook_manager.call(
            'after_test_epoch', runner, metrics=metrics)
        summary = (', '.join(f'{key}: {value}' for key, value in metrics.items())
                   if metrics else '(no evaluator configured)')
        runner.logger.info(f'Test metrics -> {summary}')
        return metrics
    finally:
        runner.hook_manager.call('after_test', runner)


def main():
    args = parse_args()
    cfg = Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)
    configure_visualization(cfg, args)

    requested_distributed = args.launcher != 'none'
    setup_runtime_environment(cfg, requested_distributed)
    distributed = init_distributed_runtime(cfg, args.launcher)
    timestamp = sync_string_from_rank0(
        datetime.now().strftime('%Y%m%d_%H%M%S'), distributed)
    cfg.device = resolve_device(str(args.device), distributed)
    prepare_test_work_dir(cfg, args.config, args.work_dir)
    barrier(distributed)

    logger = build_test_logger(cfg.work_dir, timestamp)
    logger.info(f'Starting test, config: {args.config}')
    logger.info(f'Using device: {cfg.device}')
    setup_randomness(cfg, logger)

    dataset, dataloader = build_test_data(cfg, logger)
    model, loaded_checkpoint = build_test_model(
        cfg, args.checkpoint, args.strict, logger, distributed)
    runner = build_test_context(
        cfg, logger, model, dataset, dataloader, timestamp)
    evaluator = build_test_evaluator(cfg, args.out)

    runner.hook_manager.call('before_run', runner)
    if loaded_checkpoint is not None:
        runner.hook_manager.call(
            'after_load_checkpoint', runner, checkpoint=loaded_checkpoint)
    try:
        return run_test(runner, evaluator)
    finally:
        runner.hook_manager.call('after_run', runner)


if __name__ == '__main__':
    try:
        main()
    finally:
        if (torch.distributed.is_available()
                and torch.distributed.is_initialized()):
            torch.distributed.destroy_process_group()