| 133 | |
| 134 | |
| 135 | def collect_results(result_part, size, tmpdir=None): |
| 136 | rank, world_size = get_dist_info() |
| 137 | # create a tmp dir if it is not specified |
| 138 | if tmpdir is None: |
| 139 | MAX_LEN = 512 |
| 140 | # 32 is whitespace |
| 141 | dir_tensor = torch.full((MAX_LEN, ), |
| 142 | 32, |
| 143 | dtype=torch.uint8, |
| 144 | device='cuda') |
| 145 | if rank == 0: |
| 146 | tmpdir = tempfile.mkdtemp() |
| 147 | tmpdir = torch.tensor(bytearray(tmpdir.encode()), |
| 148 | dtype=torch.uint8, |
| 149 | device='cuda') |
| 150 | dir_tensor[:len(tmpdir)] = tmpdir |
| 151 | dist.broadcast(dir_tensor, 0) |
| 152 | tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip() |
| 153 | else: |
| 154 | mmcv.mkdir_or_exist(tmpdir) |
| 155 | # dump the part result to the dir |
| 156 | mmcv.dump(result_part, osp.join(tmpdir, f'part_{rank}.pkl')) |
| 157 | dist.barrier() |
| 158 | # collect all parts |
| 159 | if rank != 0: |
| 160 | return None |
| 161 | else: |
| 162 | # load results of all parts from tmp dir |
| 163 | part_list = [] |
| 164 | for i in range(world_size): |
| 165 | part_file = osp.join(tmpdir, f'part_{i}.pkl') |
| 166 | part_list.append(mmcv.load(part_file)) |
| 167 | # sort the results |
| 168 | ordered_results = [] |
| 169 | for res in zip(*part_list): |
| 170 | ordered_results.extend(list(res)) |
| 171 | # the dataloader may pad some samples |
| 172 | ordered_results = ordered_results[:size] |
| 173 | # remove tmp dir |
| 174 | shutil.rmtree(tmpdir) |
| 175 | return ordered_results |
| 176 | |
| 177 | |
| 178 | def all_gather(data): |