mirrors
/
mmcv
mirror of https://github.com/open-mmlab/mmcv.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Add type hints for mmcv/engine/test.py (#2034) 

* Add type hints for mmcv/engine/test.py

* Add type hints for mmcv/engine/test.py

* Add type hints for mmcv/engine/test.py

* fix type hint

Co-authored-by: zhouzaida <zhouzaida@163.com>
pull/2044/head
tc_exe 2022-06-07 16:12:07 +08:00 committed by GitHub
parent 3ede2789cc
commit 4061fcdc95
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 21 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
mmcv/engine/test.py
View File

@ -4,15 +4,18 @@ import pickle
import shutil import shutil
import tempfile import tempfile
import time import time
from typing import Optional
import torch import torch
import torch.distributed as dist import torch.distributed as dist
import torch.nn as nn
from torch.utils.data import DataLoader
import mmcv import mmcv
from mmcv.runner import get_dist_info from mmcv.runner import get_dist_info
def single_gpu_test(model, data_loader): def single_gpu_test(model: nn.Module, data_loader: DataLoader) -> list:
"""Test model with a single gpu. """Test model with a single gpu.
This method tests model with a single gpu and displays test progress bar. This method tests model with a single gpu and displays test progress bar.
@ -41,7 +44,10 @@ def single_gpu_test(model, data_loader):
return results return results
def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False): def multi_gpu_test(model: nn.Module,
data_loader: DataLoader,
tmpdir: Optional[str] = None,
gpu_collect: bool = False) -> Optional[list]:
"""Test model with multiple gpus. """Test model with multiple gpus.
This method tests model with multiple gpus and collects the results This method tests model with multiple gpus and collects the results
@ -82,13 +88,15 @@ def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False):
# collect results from all ranks # collect results from all ranks
if gpu_collect: if gpu_collect:
results = collect_results_gpu(results, len(dataset)) result_from_ranks = collect_results_gpu(results, len(dataset))
else: else:
results = collect_results_cpu(results, len(dataset), tmpdir) result_from_ranks = collect_results_cpu(results, len(dataset), tmpdir)
return results return result_from_ranks
def collect_results_cpu(result_part, size, tmpdir=None): def collect_results_cpu(result_part: list,
size: int,
tmpdir: Optional[str] = None) -> Optional[list]:
"""Collect results under cpu mode. """Collect results under cpu mode.
On cpu mode, this function will save the results on different gpus to On cpu mode, this function will save the results on different gpus to
@ -126,7 +134,8 @@ def collect_results_cpu(result_part, size, tmpdir=None):
else: else:
mmcv.mkdir_or_exist(tmpdir) mmcv.mkdir_or_exist(tmpdir)
# dump the part result to the dir # dump the part result to the dir
mmcv.dump(result_part, osp.join(tmpdir, f'part_{rank}.pkl')) part_file = osp.join(tmpdir, f'part_{rank}.pkl') # type: ignore
mmcv.dump(result_part, part_file)
dist.barrier() dist.barrier()
# collect all parts # collect all parts
if rank != 0: if rank != 0:
@ -135,7 +144,7 @@ def collect_results_cpu(result_part, size, tmpdir=None):
# load results of all parts from tmp dir # load results of all parts from tmp dir
part_list = [] part_list = []
for i in range(world_size): for i in range(world_size):
part_file = osp.join(tmpdir, f'part_{i}.pkl') part_file = osp.join(tmpdir, f'part_{i}.pkl') # type: ignore
part_result = mmcv.load(part_file) part_result = mmcv.load(part_file)
# When data is severely insufficient, an empty part_result # When data is severely insufficient, an empty part_result
# on a certain gpu could makes the overall outputs empty. # on a certain gpu could makes the overall outputs empty.
@ -148,11 +157,11 @@ def collect_results_cpu(result_part, size, tmpdir=None):
# the dataloader may pad some samples # the dataloader may pad some samples
ordered_results = ordered_results[:size] ordered_results = ordered_results[:size]
# remove tmp dir # remove tmp dir
shutil.rmtree(tmpdir) shutil.rmtree(tmpdir) # type: ignore
return ordered_results return ordered_results
def collect_results_gpu(result_part, size): def collect_results_gpu(result_part: list, size: int) -> Optional[list]:
"""Collect results under gpu mode. """Collect results under gpu mode.
On gpu mode, this function will encode results to gpu tensors and use gpu On gpu mode, this function will encode results to gpu tensors and use gpu
@ -200,3 +209,5 @@ def collect_results_gpu(result_part, size):
# the dataloader may pad some samples # the dataloader may pad some samples
ordered_results = ordered_results[:size] ordered_results = ordered_results[:size]
return ordered_results return ordered_results
else:
return None