Collect results in gpu.
(result_part, size)
| 139 | |
| 140 | |
| 141 | def collect_results_gpu(result_part, size): |
| 142 | """Collect results in gpu.""" |
| 143 | rank, world_size = get_dist_info() |
| 144 | |
| 145 | # dump result part to tensor with pickle |
| 146 | part_tensor = torch.tensor(bytearray(pickle.dumps(result_part)), |
| 147 | dtype=torch.uint8, |
| 148 | device='cuda') |
| 149 | # gather all result part tensor shape |
| 150 | shape_tensor = torch.tensor(part_tensor.shape, device='cuda') |
| 151 | shape_list = [shape_tensor.clone() for _ in range(world_size)] |
| 152 | dist.all_gather(shape_list, shape_tensor) |
| 153 | # padding result part tensor to max length |
| 154 | shape_max = torch.tensor(shape_list).max() |
| 155 | part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda') |
| 156 | part_send[:shape_tensor[0]] = part_tensor |
| 157 | part_recv_list = [ |
| 158 | part_tensor.new_zeros(shape_max) for _ in range(world_size) |
| 159 | ] |
| 160 | # gather all result part |
| 161 | dist.all_gather(part_recv_list, part_send) |
| 162 | |
| 163 | if rank == 0: |
| 164 | part_list = [] |
| 165 | for recv, shape in zip(part_recv_list, shape_list): |
| 166 | part_result = pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()) |
| 167 | part_list.append(part_result) |
| 168 | # sort the results |
| 169 | ordered_results = [] |
| 170 | for res in zip(*part_list): |
| 171 | ordered_results.extend(list(res)) |
| 172 | # the dataloader may pad some samples |
| 173 | ordered_results = ordered_results[:size] |
| 174 | return ordered_results |
no test coverage detected