-
Notifications
You must be signed in to change notification settings - Fork 112
/
test.py
222 lines (191 loc) · 8.19 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import time
from typing import List, Optional
import mmcv
import torch
import torch.nn as nn
from mmcv.image import tensor2imgs
from mmcv.parallel import is_module_wrapper
from mmcv.runner import get_dist_info
from mmdet.apis.test import collect_results_cpu, collect_results_gpu
from mmdet.utils import get_root_logger
from torch.utils.data import DataLoader
def single_gpu_test(model: nn.Module,
data_loader: DataLoader,
show: bool = False,
out_dir: Optional[str] = None,
show_score_thr: float = 0.3) -> List:
"""Test model with single gpu for meta-learning based detector.
The model forward function requires `mode`, while in mmdet it requires
`return_loss`. And the `encode_mask_results` is removed.
Args:
model (nn.Module): Model to be tested.
data_loader (DataLoader): Pytorch data loader.
show (bool): Whether to show the image. Default: False.
out_dir (str | None): The directory to write the image. Default: None.
show_score_thr (float): Minimum score of bboxes to be shown.
Default: 0.3.
Returns:
list: The prediction results.
"""
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
# forward in `test` mode
result = model(mode='test', rescale=True, **data)
batch_size = len(result)
if show or out_dir:
# make sure each time only one image to be shown
if batch_size == 1 and isinstance(data['img'][0], torch.Tensor):
img_tensor = data['img'][0]
else:
img_tensor = data['img'][0].data[0]
img_metas = data['img_metas'][0].data[0]
imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg'])
assert len(imgs) == len(img_metas)
for j, (img, img_meta) in enumerate(zip(imgs, img_metas)):
h, w, _ = img_meta['img_shape']
img_show = img[:h, :w, :]
ori_h, ori_w = img_meta['ori_shape'][:-1]
img_show = mmcv.imresize(img_show, (ori_w, ori_h))
if out_dir:
out_file = osp.join(out_dir, img_meta['ori_filename'])
else:
out_file = None
if is_module_wrapper(model):
model.module.show_result(
img_show,
result[i],
show=show,
out_file=out_file,
score_thr=show_score_thr)
else:
model.show_result(
img_show,
result[i],
show=show,
out_file=out_file,
score_thr=show_score_thr)
results.extend(result)
prog_bar.update(batch_size)
return results
def multi_gpu_test(model: nn.Module,
data_loader: DataLoader,
tmpdir: str = None,
gpu_collect: bool = False) -> List:
"""Test model with multiple gpus for meta-learning based detector.
The model forward function requires `mode`, while in mmdet it requires
`return_loss`. And the `encode_mask_results` is removed.
This method tests model with multiple gpus and collects the results
under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'
it encodes results to gpu tensors and use gpu communication for results
collection. On cpu mode it saves the results on different gpus to 'tmpdir'
and collects them by the rank 0 worker.
Args:
model (nn.Module): Model to be tested.
data_loader (Dataloader): Pytorch data loader.
tmpdir (str): Path of directory to save the temporary results from
different gpus under cpu mode. Default: None.
gpu_collect (bool): Option to use either gpu or cpu to collect results.
Default: False.
Returns:
list: The prediction results.
"""
model.eval()
results = []
dataset = data_loader.dataset
rank, world_size = get_dist_info()
if rank == 0:
prog_bar = mmcv.ProgressBar(len(dataset))
time.sleep(2) # This line can prevent deadlock problem in some cases.
for i, data in enumerate(data_loader):
with torch.no_grad():
# forward in `test` mode
result = model(mode='test', rescale=True, **data)
results.extend(result)
if rank == 0:
batch_size = len(result)
prog_bar.update(batch_size * world_size)
# collect results from all ranks
if gpu_collect:
results = collect_results_gpu(results, len(dataset))
else:
results = collect_results_cpu(results, len(dataset), tmpdir)
return results
def single_gpu_model_init(model: nn.Module, data_loader: DataLoader) -> List:
"""Forward support images for meta-learning based detector initialization.
The function usually will be called before `single_gpu_test` in
`QuerySupportEvalHook`. It firstly forwards support images with
`mode=model_init` and the features will be saved in the model.
Then it will call `:func:model_init` to process the extracted features
of support images to finish the model initialization.
Args:
model (nn.Module): Model used for extracting support template features.
data_loader (nn.Dataloader): Pytorch data loader.
Returns:
list[Tensor]: Extracted support template features.
"""
model.eval()
results = []
dataset = data_loader.dataset
logger = get_root_logger()
logger.info('starting model initialization...')
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
# forward in `model_init` mode
result = model(mode='model_init', **data)
results.append(result)
prog_bar.update(num_tasks=len(data['img_metas'].data[0]))
# `model_init` will process the forward features saved in model.
if is_module_wrapper(model):
model.module.model_init()
else:
model.model_init()
logger.info('model initialization done.')
return results
def multi_gpu_model_init(model: nn.Module, data_loader: DataLoader) -> List:
"""Forward support images for meta-learning based detector initialization.
The function usually will be called before `single_gpu_test` in
`QuerySupportEvalHook`. It firstly forwards support images with
`mode=model_init` and the features will be saved in the model.
Then it will call `:func:model_init` to process the extracted features
of support images to finish the model initialization.
Noted that the `data_loader` should NOT use distributed sampler, all the
models in different gpus should be initialized with same images.
Args:
model (nn.Module): Model used for extracting support template features.
data_loader (nn.Dataloader): Pytorch data loader.
Returns:
list[Tensor]: Extracted support template features.
"""
model.eval()
results = []
dataset = data_loader.dataset
rank, _ = get_dist_info()
if rank == 0:
logger = get_root_logger()
logger.info('starting model initialization...')
prog_bar = mmcv.ProgressBar(len(dataset))
time.sleep(2) # This line can prevent deadlock problem in some cases.
# the model_init dataloader do not use distributed sampler to make sure
# all of the gpus get the same initialization
for i, data in enumerate(data_loader):
with torch.no_grad():
# forward in `model_init` mode
result = model(mode='model_init', **data)
results.append(result)
if rank == 0:
prog_bar.update(num_tasks=len(data['img_metas'].data[0]))
# model_init function will process the forward features saved in model.
if is_module_wrapper(model):
model.module.model_init()
else:
model.model_init()
if rank == 0:
logger.info('model initialization done.')
return results