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