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