Parallel testing on multiple GPUs. Args: model_cls (type): Model class type. model_kwargs (dict): Arguments to init the model. checkpoint (str): Checkpoint filepath. dataset (:obj:`Dataset`): The dataset to be tested. data_func (callable): The function th
(model_cls,
model_kwargs,
checkpoint,
dataset,
data_func,
gpus,
workers_per_gpu=1)
| 22 | |
| 23 | |
| 24 | def parallel_test(model_cls, |
| 25 | model_kwargs, |
| 26 | checkpoint, |
| 27 | dataset, |
| 28 | data_func, |
| 29 | gpus, |
| 30 | workers_per_gpu=1): |
| 31 | """Parallel testing on multiple GPUs. |
| 32 | |
| 33 | Args: |
| 34 | model_cls (type): Model class type. |
| 35 | model_kwargs (dict): Arguments to init the model. |
| 36 | checkpoint (str): Checkpoint filepath. |
| 37 | dataset (:obj:`Dataset`): The dataset to be tested. |
| 38 | data_func (callable): The function that generates model inputs. |
| 39 | gpus (list[int]): GPU ids to be used. |
| 40 | workers_per_gpu (int): Number of processes on each GPU. It is possible |
| 41 | to run multiple workers on each GPU. |
| 42 | |
| 43 | Returns: |
| 44 | list: Test results. |
| 45 | """ |
| 46 | ctx = multiprocessing.get_context('spawn') |
| 47 | idx_queue = ctx.Queue() |
| 48 | result_queue = ctx.Queue() |
| 49 | num_workers = len(gpus) * workers_per_gpu |
| 50 | workers = [ |
| 51 | ctx.Process( |
| 52 | target=worker_func, |
| 53 | args=(model_cls, model_kwargs, checkpoint, dataset, data_func, |
| 54 | gpus[i % len(gpus)], idx_queue, result_queue)) |
| 55 | for i in range(num_workers) |
| 56 | ] |
| 57 | for w in workers: |
| 58 | w.daemon = True |
| 59 | w.start() |
| 60 | |
| 61 | for i in range(len(dataset)): |
| 62 | idx_queue.put(i) |
| 63 | |
| 64 | results = [None for _ in range(len(dataset))] |
| 65 | prog_bar = mmcv.ProgressBar(task_num=len(dataset)) |
| 66 | for _ in range(len(dataset)): |
| 67 | idx, res = result_queue.get() |
| 68 | results[idx] = res |
| 69 | prog_bar.update() |
| 70 | print('\n') |
| 71 | for worker in workers: |
| 72 | worker.terminate() |
| 73 | |
| 74 | return results |