Collect results in cpu.
(result_part, size, tmpdir=None)
| 91 | |
| 92 | |
| 93 | def collect_results_cpu(result_part, size, tmpdir=None): |
| 94 | """Collect results in cpu.""" |
| 95 | rank, world_size = get_dist_info() |
| 96 | # create a tmp dir if it is not specified |
| 97 | if tmpdir is None: |
| 98 | MAX_LEN = 512 |
| 99 | # 32 is whitespace |
| 100 | dir_tensor = torch.full((MAX_LEN, ), |
| 101 | 32, |
| 102 | dtype=torch.uint8, |
| 103 | device='cuda') |
| 104 | if rank == 0: |
| 105 | mmcv.mkdir_or_exist('.dist_test') |
| 106 | tmpdir = tempfile.mkdtemp(dir='.dist_test') |
| 107 | tmpdir = torch.tensor(bytearray(tmpdir.encode()), |
| 108 | dtype=torch.uint8, |
| 109 | device='cuda') |
| 110 | dir_tensor[:len(tmpdir)] = tmpdir |
| 111 | dist.broadcast(dir_tensor, 0) |
| 112 | tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip() |
| 113 | else: |
| 114 | mmcv.mkdir_or_exist(tmpdir) |
| 115 | # dump the part result to the dir |
| 116 | mmcv.dump(result_part, osp.join(tmpdir, f'part_{rank}.pkl')) |
| 117 | dist.barrier() |
| 118 | # collect all parts |
| 119 | if rank != 0: |
| 120 | return None |
| 121 | else: |
| 122 | # load results of all parts from tmp dir |
| 123 | part_list = [] |
| 124 | for i in range(world_size): |
| 125 | part_file = osp.join(tmpdir, f'part_{i}.pkl') |
| 126 | part_result = mmcv.load(part_file) |
| 127 | part_list.append(part_result) |
| 128 | |
| 129 | # sort the results |
| 130 | ordered_results = [] |
| 131 | for res in zip(*part_list): |
| 132 | ordered_results.extend(list(res)) |
| 133 | # the dataloader may pad some samples |
| 134 | |
| 135 | ordered_results = ordered_results[:size] |
| 136 | # remove tmp dir |
| 137 | shutil.rmtree(tmpdir) |
| 138 | return ordered_results |
| 139 | |
| 140 | |
| 141 | def collect_results_gpu(result_part, size): |