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
(model, data_loader, tmpdir=None, gpu_collect=False)
| 43 | |
| 44 | def custom_multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False, occ_threshold=0.25): |
| 45 | """Test model with multiple gpus. |
| 46 | This method tests model with multiple gpus and collects the results |
| 47 | under two different modes: gpu and cpu modes. By setting 'gpu_collect=True' |
| 48 | it encodes results to gpu tensors and use gpu communication for results |
| 49 | collection. On cpu mode it saves the results on different gpus to 'tmpdir' |
| 50 | and collects them by the rank 0 worker. |
| 51 | Args: |
| 52 | model (nn.Module): Model to be tested. |
| 53 | data_loader (nn.Dataloader): Pytorch data loader. |
| 54 | tmpdir (str): Path of directory to save the temporary results from |
| 55 | different gpus under cpu mode. |
| 56 | gpu_collect (bool): Option to use either gpu or cpu to collect results. |
| 57 | Returns: |
| 58 | list: The prediction results. |
| 59 | """ |
| 60 | model.eval() |
| 61 | bbox_results = [] |
| 62 | mask_results = [] |
| 63 | occupancy_results = [] |
| 64 | flow_results = [] |
| 65 | dataset = data_loader.dataset |
| 66 | rank, world_size = get_dist_info() |
| 67 | if rank == 0: |
| 68 | prog_bar = mmcv.ProgressBar(len(dataset)) |
| 69 | time.sleep(2) # This line can prevent deadlock problem in some cases. |
| 70 | have_mask = False |
| 71 | for i, data in enumerate(data_loader): |
| 72 | with torch.no_grad(): |
| 73 | result, occ_results = model(return_loss=False, rescale=True, occ_threshold=occ_threshold, **data) |
| 74 | if isinstance(result, dict): |
| 75 | if 'bbox_results' in result.keys(): |
| 76 | bbox_result = result['bbox_results'] |
| 77 | batch_size = len(result['bbox_results']) |
| 78 | bbox_results.extend(bbox_result) |
| 79 | if 'mask_results' in result.keys() and result['mask_results'] is not None: |
| 80 | mask_result = custom_encode_mask_results(result['mask_results']) |
| 81 | mask_results.extend(mask_result) |
| 82 | have_mask = True |
| 83 | elif result is None: |
| 84 | bbox_results = [] |
| 85 | else: |
| 86 | batch_size = len(result) |
| 87 | bbox_results.extend(result) |
| 88 | |
| 89 | occupancy_preds = occ_results['occupancy_preds'] |
| 90 | flow_preds = occ_results['flow_preds'] |
| 91 | |
| 92 | if occupancy_preds is not None: |
| 93 | batch_size = 1 |
| 94 | occupancy_results.extend([occupancy_preds]) |
| 95 | if flow_preds is not None: |
| 96 | flow_results.extend([flow_preds]) |
no test coverage detected