Function do_test

scripts/train.py:197–274

(cfg, model, is_last=False, use_tta=False)

def do_test(cfg, model, is_last=False, use_tta=False):
    if not cfg.TEST.ENABLED:
        LOG.warning("Test is disabled.")
        return {}

    dataset_names = [cfg.DATASETS.TEST.NAME]  # NOTE: only support single test dataset for now.

    if use_tta:
        LOG.info("Starting inference with test-time augmentation.")
        if isinstance(model, DistributedDataParallel):
            model.module.postprocess_in_inference = False
        else:
            model.postprocess_in_inference = False
        model = build_tta_model(cfg, model)

    test_results = OrderedDict()
    for dataset_name in dataset_names:
        # output directory for this dataset.
        dset_output_dir = get_inference_output_dir(dataset_name, is_last=is_last, use_tta=use_tta)

        # What evaluators are used for this dataset?
        evaluator_names = MetadataCatalog.get(dataset_name).evaluators
        evaluators = []
        for evaluator_name in evaluator_names:
            evaluator = get_evaluator(cfg, dataset_name, evaluator_name, dset_output_dir)
            evaluators.append(evaluator)
        evaluator = DatasetEvaluators(evaluators)

        mapper = get_dataset_mapper(cfg, is_train=False)
        dataloader, dataset_dicts = build_test_dataloader(cfg, dataset_name, mapper)

        per_dataset_results = inference_on_dataset(model, dataloader, evaluator)
        if use_tta:
            per_dataset_results = OrderedDict({k + '-tta': v for k, v in per_dataset_results.items()})
        test_results[dataset_name] = per_dataset_results

        if cfg.VIS.PREDICTIONS_ENABLED and d2_comm.is_main_process():
            visualizer_names = MetadataCatalog.get(dataset_name).pred_visualizers
            # Randomly (but deterministically) select what samples to visualize.
            # The samples are shared across all visualizers and iterations.
            sampled_dataset_dicts, inds = random_sample_dataset_dicts(
                dataset_name, num_samples=cfg.VIS.PREDICTIONS_MAX_NUM_SAMPLES
            )

            viz_images = defaultdict(dict)
            for viz_name in visualizer_names:
                LOG.info(f"Running prediction visualizer: {viz_name}")
                visualizer = get_predictions_visualizer(cfg, viz_name, dataset_name, dset_output_dir)
                for x in tqdm(sampled_dataset_dicts):
                    sample_id = x['sample_id']
                    viz_images[sample_id].update(visualizer.visualize(x))

            save_vis(viz_images, dset_output_dir, "visualization")

            if cfg.WANDB.ENABLED:
                LOG.info(f"Uploading prediction visualization to W&B: {dataset_name}")
                for sample_id in viz_images.keys():
                    viz_images[sample_id] = mosaic(list(viz_images[sample_id].values()))