Function multi_gpu_test

easycv/apis/test.py:155–249

Test model with multiple gpus. This method tests model with multiple gpus and collects the results under two different modes: gpu and cpu modes. By setting 'gpu_collect=True' it encodes results to gpu tensors and use gpu communication for results collection. On cpu mode it saves the

(model,
                   data_loader,
                   mode='test',
                   tmpdir=None,
                   gpu_collect=False,
                   use_fp16=False,
                   **kwargs)

def multi_gpu_test(model,
                   data_loader,
                   mode='test',
                   tmpdir=None,
                   gpu_collect=False,
                   use_fp16=False,
                   **kwargs):
    """Test model with multiple gpus.

    This method tests model with multiple gpus and collects the results
    under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'
    it encodes results to gpu tensors and use gpu communication for results
    collection. On cpu mode it saves the results on different gpus to 'tmpdir'
    and collects them by the rank 0 worker.

    Args:
        model (nn.Module): Model to be tested.
        data_loader (nn.Dataloader): Pytorch data loader.
        model (str): mode for model to forward
        tmpdir (str): Path of directory to save the temporary results from
            different gpus under cpu mode.
        gpu_collect (bool): Option to use either gpu or cpu to collect results.
        use_fp16:  Use fp16 inference

    Returns:
        list: The prediction results.
    """
    if use_fp16:
        device = next(model.parameters()).device
        assert device.type == 'cuda', 'fp16 can only be used in gpu, model is placed on cpu'
        model.half()

    model.eval()
    results = {}
    rank, world_size = get_dist_info()

    if hasattr(data_loader, 'dataset'):  # normal dataloader
        data_len = len(data_loader.dataset)
    else:
        data_len = len(data_loader) * data_loader.batch_size * world_size

    if rank == 0:
        prog_bar = mmcv.ProgressBar(data_len)
    time.sleep(2)  # This line can prevent deadlock problem in some cases.

    for i, data in enumerate(data_loader):
        with torch.no_grad():
            result = model(**data, mode=mode)
            # # encode mask results
            # if isinstance(result, tuple):
            #     bbox_results, mask_results = result
            #     encoded_mask_results = encode_mask_results(mask_results)
            #     result = bbox_results, encoded_mask_results

        for k, v in result.items():
            if k not in results:
                results[k] = []