Spaces:
Running
on
T4
Running
on
T4
# Copyright (c) OpenMMLab. All rights reserved. | |
import os | |
import random | |
import numpy as np | |
import pytest | |
import torch | |
from mmengine import Config | |
try: | |
import importlib | |
importlib.import_module('mmdeploy') | |
except ImportError: | |
pytest.skip('mmdeploy is not installed.', allow_module_level=True) | |
from mmdeploy.codebase import import_codebase | |
from mmdeploy.utils import Backend | |
from mmdeploy.utils.config_utils import register_codebase | |
from mmdeploy.utils.test import (WrapModel, check_backend, get_model_outputs, | |
get_rewrite_outputs) | |
try: | |
codebase = register_codebase('mmyolo') | |
import_codebase(codebase, ['mmyolo.deploy']) | |
except ImportError: | |
pytest.skip('mmyolo is not installed.', allow_module_level=True) | |
def seed_everything(seed=1029): | |
random.seed(seed) | |
os.environ['PYTHONHASHSEED'] = str(seed) | |
np.random.seed(seed) | |
torch.manual_seed(seed) | |
if torch.cuda.is_available(): | |
torch.cuda.manual_seed(seed) | |
torch.cuda.manual_seed_all(seed) # if you are using multi-GPU. | |
torch.backends.cudnn.benchmark = False | |
torch.backends.cudnn.deterministic = True | |
torch.backends.cudnn.enabled = False | |
def get_yolov5_head_model(): | |
"""YOLOv5 Head Config.""" | |
test_cfg = Config( | |
dict( | |
multi_label=True, | |
nms_pre=30000, | |
score_thr=0.001, | |
nms=dict(type='nms', iou_threshold=0.65), | |
max_per_img=300)) | |
from mmyolo.models.dense_heads import YOLOv5Head | |
head_module = dict( | |
type='YOLOv5HeadModule', | |
num_classes=4, | |
in_channels=[2, 4, 8], | |
featmap_strides=[8, 16, 32], | |
num_base_priors=1) | |
model = YOLOv5Head(head_module, test_cfg=test_cfg) | |
model.requires_grad_(False) | |
return model | |
def test_yolov5_head_predict_by_feat(backend_type: Backend): | |
"""Test predict_by_feat rewrite of YOLOXHead.""" | |
check_backend(backend_type) | |
yolov5_head = get_yolov5_head_model() | |
yolov5_head.cpu().eval() | |
s = 256 | |
batch_img_metas = [{ | |
'scale_factor': (1.0, 1.0), | |
'pad_shape': (s, s, 3), | |
'img_shape': (s, s, 3), | |
'ori_shape': (s, s, 3) | |
}] | |
output_names = ['dets', 'labels'] | |
deploy_cfg = Config( | |
dict( | |
backend_config=dict(type=backend_type.value), | |
onnx_config=dict(output_names=output_names, input_shape=None), | |
codebase_config=dict( | |
type='mmyolo', | |
task='ObjectDetection', | |
post_processing=dict( | |
score_threshold=0.05, | |
iou_threshold=0.5, | |
max_output_boxes_per_class=20, | |
pre_top_k=-1, | |
keep_top_k=10, | |
background_label_id=-1, | |
), | |
module=['mmyolo.deploy']))) | |
seed_everything(1234) | |
cls_scores = [ | |
torch.rand(1, yolov5_head.num_classes * yolov5_head.num_base_priors, | |
4 * pow(2, i), 4 * pow(2, i)) for i in range(3, 0, -1) | |
] | |
seed_everything(5678) | |
bbox_preds = [ | |
torch.rand(1, 4 * yolov5_head.num_base_priors, 4 * pow(2, i), | |
4 * pow(2, i)) for i in range(3, 0, -1) | |
] | |
seed_everything(9101) | |
objectnesses = [ | |
torch.rand(1, 1 * yolov5_head.num_base_priors, 4 * pow(2, i), | |
4 * pow(2, i)) for i in range(3, 0, -1) | |
] | |
# to get outputs of pytorch model | |
model_inputs = { | |
'cls_scores': cls_scores, | |
'bbox_preds': bbox_preds, | |
'objectnesses': objectnesses, | |
'batch_img_metas': batch_img_metas, | |
'with_nms': True | |
} | |
model_outputs = get_model_outputs(yolov5_head, 'predict_by_feat', | |
model_inputs) | |
# to get outputs of onnx model after rewrite | |
wrapped_model = WrapModel( | |
yolov5_head, | |
'predict_by_feat', | |
batch_img_metas=batch_img_metas, | |
with_nms=True) | |
rewrite_inputs = { | |
'cls_scores': cls_scores, | |
'bbox_preds': bbox_preds, | |
'objectnesses': objectnesses, | |
} | |
rewrite_outputs, is_backend_output = get_rewrite_outputs( | |
wrapped_model=wrapped_model, | |
model_inputs=rewrite_inputs, | |
deploy_cfg=deploy_cfg) | |
if is_backend_output: | |
# hard code to make two tensors with the same shape | |
# rewrite and original codes applied different nms strategy | |
min_shape = min(model_outputs[0].bboxes.shape[0], | |
rewrite_outputs[0].shape[1], 5) | |
for i in range(len(model_outputs)): | |
rewrite_outputs[0][i, :min_shape, 0::2] = \ | |
rewrite_outputs[0][i, :min_shape, 0::2].clamp_(0, s) | |
rewrite_outputs[0][i, :min_shape, 1::2] = \ | |
rewrite_outputs[0][i, :min_shape, 1::2].clamp_(0, s) | |
assert np.allclose( | |
model_outputs[i].bboxes[:min_shape], | |
rewrite_outputs[0][i, :min_shape, :4], | |
rtol=1e-03, | |
atol=1e-05) | |
assert np.allclose( | |
model_outputs[i].scores[:min_shape], | |
rewrite_outputs[0][i, :min_shape, 4], | |
rtol=1e-03, | |
atol=1e-05) | |
assert np.allclose( | |
model_outputs[i].labels[:min_shape], | |
rewrite_outputs[1][i, :min_shape], | |
rtol=1e-03, | |
atol=1e-05) | |
else: | |
assert rewrite_outputs is not None | |