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