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, tmpdir=None, gpu_collect=False)
| 33 | |
| 34 | |
| 35 | def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False): |
| 36 | """Test model with multiple gpus. |
| 37 | This method tests model with multiple gpus and collects the results |
| 38 | under two different modes: gpu and cpu modes. By setting 'gpu_collect=True' |
| 39 | it encodes results to gpu tensors and use gpu communication for results |
| 40 | collection. On cpu mode it saves the results on different gpus to 'tmpdir' |
| 41 | and collects them by the rank 0 worker. |
| 42 | Args: |
| 43 | model (nn.Module): Model to be tested. |
| 44 | data_loader (nn.Dataloader): Pytorch data loader. |
| 45 | tmpdir (str): Path of directory to save the temporary results from |
| 46 | different gpus under cpu mode. |
| 47 | gpu_collect (bool): Option to use either gpu or cpu to collect results. |
| 48 | Returns: |
| 49 | list: The prediction results. |
| 50 | """ |
| 51 | model.eval() |
| 52 | results = [] |
| 53 | dataset = data_loader.dataset |
| 54 | rank, world_size = get_dist_info() |
| 55 | if rank == 0: |
| 56 | # Check if tmpdir is valid for cpu_collect |
| 57 | if (not gpu_collect) and (tmpdir is not None and osp.exists(tmpdir)): |
| 58 | raise OSError((f'The tmpdir {tmpdir} already exists.', |
| 59 | ' Since tmpdir will be deleted after testing,', |
| 60 | ' please make sure you specify an empty one.')) |
| 61 | prog_bar = mmcv.ProgressBar(len(dataset)) |
| 62 | time.sleep(2) # This line can prevent deadlock problem in some cases. |
| 63 | for i, data in enumerate(data_loader): |
| 64 | with torch.no_grad(): |
| 65 | result = model(return_loss=False, **data) |
| 66 | if isinstance(result, list): |
| 67 | results.extend(result) |
| 68 | else: |
| 69 | results.append(result) |
| 70 | |
| 71 | if rank == 0: |
| 72 | batch_size = data['motion'].size(0) |
| 73 | for _ in range(batch_size * world_size): |
| 74 | prog_bar.update() |
| 75 | |
| 76 | # collect results from all ranks |
| 77 | if gpu_collect: |
| 78 | results = collect_results_gpu(results, len(dataset)) |
| 79 | else: |
| 80 | results = collect_results_cpu(results, len(dataset), tmpdir) |
| 81 | return results |
| 82 | |
| 83 | |
| 84 | def collect_results_cpu(result_part, size, tmpdir=None): |
no test coverage detected