454 lines
14 KiB
Python
454 lines
14 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import os.path as osp
|
|
from functools import partial
|
|
|
|
import mmcv
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from mmcv.cnn import Scale
|
|
|
|
from mmdet import digit_version
|
|
from mmdet.models import build_detector
|
|
from mmdet.models.dense_heads import (FCOSHead, FSAFHead, RetinaHead, SSDHead,
|
|
YOLOV3Head)
|
|
from .utils import ort_validate
|
|
|
|
data_path = osp.join(osp.dirname(__file__), 'data')
|
|
|
|
if digit_version(torch.__version__) <= digit_version('1.5.0'):
|
|
pytest.skip(
|
|
'ort backend does not support version below 1.5.0',
|
|
allow_module_level=True)
|
|
|
|
|
|
def test_cascade_onnx_export():
|
|
|
|
config_path = './configs/cascade_rcnn/cascade_rcnn_r50_fpn_1x_coco.py'
|
|
cfg = mmcv.Config.fromfile(config_path)
|
|
model = build_detector(cfg.model, test_cfg=cfg.get('test_cfg'))
|
|
with torch.no_grad():
|
|
model.forward = partial(model.forward, img_metas=[[dict()]])
|
|
|
|
dynamic_axes = {
|
|
'input_img': {
|
|
0: 'batch',
|
|
2: 'width',
|
|
3: 'height'
|
|
},
|
|
'dets': {
|
|
0: 'batch',
|
|
1: 'num_dets',
|
|
},
|
|
'labels': {
|
|
0: 'batch',
|
|
1: 'num_dets',
|
|
},
|
|
}
|
|
torch.onnx.export(
|
|
model, [torch.rand(1, 3, 400, 500)],
|
|
'tmp.onnx',
|
|
output_names=['dets', 'labels'],
|
|
input_names=['input_img'],
|
|
keep_initializers_as_inputs=True,
|
|
do_constant_folding=True,
|
|
verbose=False,
|
|
opset_version=11,
|
|
dynamic_axes=dynamic_axes)
|
|
|
|
|
|
def test_faster_onnx_export():
|
|
|
|
config_path = './configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py'
|
|
cfg = mmcv.Config.fromfile(config_path)
|
|
model = build_detector(cfg.model, test_cfg=cfg.get('test_cfg'))
|
|
with torch.no_grad():
|
|
model.forward = partial(model.forward, img_metas=[[dict()]])
|
|
|
|
dynamic_axes = {
|
|
'input_img': {
|
|
0: 'batch',
|
|
2: 'width',
|
|
3: 'height'
|
|
},
|
|
'dets': {
|
|
0: 'batch',
|
|
1: 'num_dets',
|
|
},
|
|
'labels': {
|
|
0: 'batch',
|
|
1: 'num_dets',
|
|
},
|
|
}
|
|
torch.onnx.export(
|
|
model, [torch.rand(1, 3, 400, 500)],
|
|
'tmp.onnx',
|
|
output_names=['dets', 'labels'],
|
|
input_names=['input_img'],
|
|
keep_initializers_as_inputs=True,
|
|
do_constant_folding=True,
|
|
verbose=False,
|
|
opset_version=11,
|
|
dynamic_axes=dynamic_axes)
|
|
|
|
|
|
def retinanet_config():
|
|
"""RetinanNet Head Config."""
|
|
head_cfg = dict(
|
|
stacked_convs=6,
|
|
feat_channels=2,
|
|
anchor_generator=dict(
|
|
type='AnchorGenerator',
|
|
octave_base_scale=4,
|
|
scales_per_octave=3,
|
|
ratios=[0.5, 1.0, 2.0],
|
|
strides=[8, 16, 32, 64, 128]),
|
|
bbox_coder=dict(
|
|
type='DeltaXYWHBBoxCoder',
|
|
target_means=[.0, .0, .0, .0],
|
|
target_stds=[1.0, 1.0, 1.0, 1.0]))
|
|
|
|
test_cfg = mmcv.Config(
|
|
dict(
|
|
deploy_nms_pre=0,
|
|
min_bbox_size=0,
|
|
score_thr=0.05,
|
|
nms=dict(type='nms', iou_threshold=0.5),
|
|
max_per_img=100))
|
|
|
|
model = RetinaHead(
|
|
num_classes=4, in_channels=1, test_cfg=test_cfg, **head_cfg)
|
|
model.requires_grad_(False)
|
|
|
|
return model
|
|
|
|
|
|
def test_retina_head_forward_single():
|
|
"""Test RetinaNet Head single forward in torch and onnxruntime env."""
|
|
retina_model = retinanet_config()
|
|
|
|
feat = torch.rand(1, retina_model.in_channels, 32, 32)
|
|
# validate the result between the torch and ort
|
|
ort_validate(retina_model.forward_single, feat)
|
|
|
|
|
|
def test_retina_head_forward():
|
|
"""Test RetinaNet Head forward in torch and onnxruntime env."""
|
|
retina_model = retinanet_config()
|
|
s = 128
|
|
# RetinaNet head expects a multiple levels of features per image
|
|
feats = [
|
|
torch.rand(1, retina_model.in_channels, s // (2**(i + 2)),
|
|
s // (2**(i + 2))) # [32, 16, 8, 4, 2]
|
|
for i in range(len(retina_model.prior_generator.strides))
|
|
]
|
|
ort_validate(retina_model.forward, feats)
|
|
|
|
|
|
def test_retinanet_head_onnx_export():
|
|
"""Test RetinaNet Head _get_bboxes() in torch and onnxruntime env."""
|
|
retina_model = retinanet_config()
|
|
s = 128
|
|
img_metas = [{
|
|
'img_shape_for_onnx': torch.Tensor([s, s]),
|
|
'scale_factor': np.ones(4),
|
|
'pad_shape': (s, s, 3),
|
|
'img_shape': (s, s, 2)
|
|
}]
|
|
|
|
# The data of retina_head_get_bboxes.pkl contains two parts:
|
|
# cls_score(list(Tensor)) and bboxes(list(Tensor)),
|
|
# where each torch.Tensor is generated by torch.rand().
|
|
# the cls_score's size: (1, 36, 32, 32), (1, 36, 16, 16),
|
|
# (1, 36, 8, 8), (1, 36, 4, 4), (1, 36, 2, 2).
|
|
# the bboxes's size: (1, 36, 32, 32), (1, 36, 16, 16),
|
|
# (1, 36, 8, 8), (1, 36, 4, 4), (1, 36, 2, 2)
|
|
retina_head_data = 'retina_head_get_bboxes.pkl'
|
|
feats = mmcv.load(osp.join(data_path, retina_head_data))
|
|
cls_score = feats[:5]
|
|
bboxes = feats[5:]
|
|
|
|
retina_model.onnx_export = partial(
|
|
retina_model.onnx_export, img_metas=img_metas, with_nms=False)
|
|
ort_validate(retina_model.onnx_export, (cls_score, bboxes))
|
|
|
|
|
|
def yolo_config():
|
|
"""YoloV3 Head Config."""
|
|
head_cfg = dict(
|
|
anchor_generator=dict(
|
|
type='YOLOAnchorGenerator',
|
|
base_sizes=[[(116, 90), (156, 198), (373, 326)],
|
|
[(30, 61), (62, 45), (59, 119)],
|
|
[(10, 13), (16, 30), (33, 23)]],
|
|
strides=[32, 16, 8]),
|
|
bbox_coder=dict(type='YOLOBBoxCoder'))
|
|
|
|
test_cfg = mmcv.Config(
|
|
dict(
|
|
deploy_nms_pre=0,
|
|
min_bbox_size=0,
|
|
score_thr=0.05,
|
|
conf_thr=0.005,
|
|
nms=dict(type='nms', iou_threshold=0.45),
|
|
max_per_img=100))
|
|
|
|
model = YOLOV3Head(
|
|
num_classes=4,
|
|
in_channels=[1, 1, 1],
|
|
out_channels=[16, 8, 4],
|
|
test_cfg=test_cfg,
|
|
**head_cfg)
|
|
model.requires_grad_(False)
|
|
# yolov3 need eval()
|
|
model.cpu().eval()
|
|
return model
|
|
|
|
|
|
def test_yolov3_head_forward():
|
|
"""Test Yolov3 head forward() in torch and ort env."""
|
|
yolo_model = yolo_config()
|
|
|
|
# Yolov3 head expects a multiple levels of features per image
|
|
feats = [
|
|
torch.rand(1, 1, 64 // (2**(i + 2)), 64 // (2**(i + 2)))
|
|
for i in range(len(yolo_model.in_channels))
|
|
]
|
|
ort_validate(yolo_model.forward, feats)
|
|
|
|
|
|
def test_yolov3_head_onnx_export():
|
|
"""Test yolov3 head get_bboxes() in torch and ort env."""
|
|
yolo_model = yolo_config()
|
|
s = 128
|
|
img_metas = [{
|
|
'img_shape_for_onnx': torch.Tensor([s, s]),
|
|
'img_shape': (s, s, 3),
|
|
'scale_factor': np.ones(4),
|
|
'pad_shape': (s, s, 3)
|
|
}]
|
|
|
|
# The data of yolov3_head_get_bboxes.pkl contains
|
|
# a list of torch.Tensor, where each torch.Tensor
|
|
# is generated by torch.rand and each tensor size is:
|
|
# (1, 27, 32, 32), (1, 27, 16, 16), (1, 27, 8, 8).
|
|
yolo_head_data = 'yolov3_head_get_bboxes.pkl'
|
|
pred_maps = mmcv.load(osp.join(data_path, yolo_head_data))
|
|
|
|
yolo_model.onnx_export = partial(
|
|
yolo_model.onnx_export, img_metas=img_metas, with_nms=False)
|
|
ort_validate(yolo_model.onnx_export, pred_maps)
|
|
|
|
|
|
def fcos_config():
|
|
"""FCOS Head Config."""
|
|
test_cfg = mmcv.Config(
|
|
dict(
|
|
deploy_nms_pre=0,
|
|
min_bbox_size=0,
|
|
score_thr=0.05,
|
|
nms=dict(type='nms', iou_threshold=0.5),
|
|
max_per_img=100))
|
|
|
|
model = FCOSHead(num_classes=4, in_channels=1, test_cfg=test_cfg)
|
|
|
|
model.requires_grad_(False)
|
|
return model
|
|
|
|
|
|
def test_fcos_head_forward_single():
|
|
"""Test fcos forward single in torch and ort env."""
|
|
fcos_model = fcos_config()
|
|
|
|
feat = torch.rand(1, fcos_model.in_channels, 32, 32)
|
|
fcos_model.forward_single = partial(
|
|
fcos_model.forward_single,
|
|
scale=Scale(1.0).requires_grad_(False),
|
|
stride=(4, ))
|
|
ort_validate(fcos_model.forward_single, feat)
|
|
|
|
|
|
def test_fcos_head_forward():
|
|
"""Test fcos forward in mutil-level feature map."""
|
|
fcos_model = fcos_config()
|
|
s = 128
|
|
feats = [
|
|
torch.rand(1, 1, s // feat_size, s // feat_size)
|
|
for feat_size in [4, 8, 16, 32, 64]
|
|
]
|
|
ort_validate(fcos_model.forward, feats)
|
|
|
|
|
|
def test_fcos_head_onnx_export():
|
|
"""Test fcos head get_bboxes() in ort."""
|
|
fcos_model = fcos_config()
|
|
s = 128
|
|
img_metas = [{
|
|
'img_shape_for_onnx': torch.Tensor([s, s]),
|
|
'img_shape': (s, s, 3),
|
|
'scale_factor': np.ones(4),
|
|
'pad_shape': (s, s, 3)
|
|
}]
|
|
|
|
cls_scores = [
|
|
torch.rand(1, fcos_model.num_classes, s // feat_size, s // feat_size)
|
|
for feat_size in [4, 8, 16, 32, 64]
|
|
]
|
|
bboxes = [
|
|
torch.rand(1, 4, s // feat_size, s // feat_size)
|
|
for feat_size in [4, 8, 16, 32, 64]
|
|
]
|
|
centerness = [
|
|
torch.rand(1, 1, s // feat_size, s // feat_size)
|
|
for feat_size in [4, 8, 16, 32, 64]
|
|
]
|
|
|
|
fcos_model.onnx_export = partial(
|
|
fcos_model.onnx_export, img_metas=img_metas, with_nms=False)
|
|
ort_validate(fcos_model.onnx_export, (cls_scores, bboxes, centerness))
|
|
|
|
|
|
def fsaf_config():
|
|
"""FSAF Head Config."""
|
|
cfg = dict(
|
|
anchor_generator=dict(
|
|
type='AnchorGenerator',
|
|
octave_base_scale=1,
|
|
scales_per_octave=1,
|
|
ratios=[1.0],
|
|
strides=[8, 16, 32, 64, 128]))
|
|
|
|
test_cfg = mmcv.Config(
|
|
dict(
|
|
deploy_nms_pre=0,
|
|
min_bbox_size=0,
|
|
score_thr=0.05,
|
|
nms=dict(type='nms', iou_threshold=0.5),
|
|
max_per_img=100))
|
|
|
|
model = FSAFHead(num_classes=4, in_channels=1, test_cfg=test_cfg, **cfg)
|
|
model.requires_grad_(False)
|
|
return model
|
|
|
|
|
|
def test_fsaf_head_forward_single():
|
|
"""Test RetinaNet Head forward_single() in torch and onnxruntime env."""
|
|
fsaf_model = fsaf_config()
|
|
|
|
feat = torch.rand(1, fsaf_model.in_channels, 32, 32)
|
|
ort_validate(fsaf_model.forward_single, feat)
|
|
|
|
|
|
def test_fsaf_head_forward():
|
|
"""Test RetinaNet Head forward in torch and onnxruntime env."""
|
|
fsaf_model = fsaf_config()
|
|
s = 128
|
|
feats = [
|
|
torch.rand(1, fsaf_model.in_channels, s // (2**(i + 2)),
|
|
s // (2**(i + 2)))
|
|
for i in range(len(fsaf_model.anchor_generator.strides))
|
|
]
|
|
ort_validate(fsaf_model.forward, feats)
|
|
|
|
|
|
def test_fsaf_head_onnx_export():
|
|
"""Test RetinaNet Head get_bboxes in torch and onnxruntime env."""
|
|
fsaf_model = fsaf_config()
|
|
s = 256
|
|
img_metas = [{
|
|
'img_shape_for_onnx': torch.Tensor([s, s]),
|
|
'scale_factor': np.ones(4),
|
|
'pad_shape': (s, s, 3),
|
|
'img_shape': (s, s, 2)
|
|
}]
|
|
|
|
# The data of fsaf_head_get_bboxes.pkl contains two parts:
|
|
# cls_score(list(Tensor)) and bboxes(list(Tensor)),
|
|
# where each torch.Tensor is generated by torch.rand().
|
|
# the cls_score's size: (1, 4, 64, 64), (1, 4, 32, 32),
|
|
# (1, 4, 16, 16), (1, 4, 8, 8), (1, 4, 4, 4).
|
|
# the bboxes's size: (1, 4, 64, 64), (1, 4, 32, 32),
|
|
# (1, 4, 16, 16), (1, 4, 8, 8), (1, 4, 4, 4).
|
|
fsaf_head_data = 'fsaf_head_get_bboxes.pkl'
|
|
feats = mmcv.load(osp.join(data_path, fsaf_head_data))
|
|
cls_score = feats[:5]
|
|
bboxes = feats[5:]
|
|
|
|
fsaf_model.onnx_export = partial(
|
|
fsaf_model.onnx_export, img_metas=img_metas, with_nms=False)
|
|
ort_validate(fsaf_model.onnx_export, (cls_score, bboxes))
|
|
|
|
|
|
def ssd_config():
|
|
"""SSD Head Config."""
|
|
cfg = dict(
|
|
anchor_generator=dict(
|
|
type='SSDAnchorGenerator',
|
|
scale_major=False,
|
|
input_size=300,
|
|
basesize_ratio_range=(0.15, 0.9),
|
|
strides=[8, 16, 32, 64, 100, 300],
|
|
ratios=[[2], [2, 3], [2, 3], [2, 3], [2], [2]]),
|
|
bbox_coder=dict(
|
|
type='DeltaXYWHBBoxCoder',
|
|
target_means=[.0, .0, .0, .0],
|
|
target_stds=[0.1, 0.1, 0.2, 0.2]))
|
|
|
|
test_cfg = mmcv.Config(
|
|
dict(
|
|
deploy_nms_pre=0,
|
|
nms=dict(type='nms', iou_threshold=0.45),
|
|
min_bbox_size=0,
|
|
score_thr=0.02,
|
|
max_per_img=200))
|
|
|
|
model = SSDHead(
|
|
num_classes=4,
|
|
in_channels=(4, 8, 4, 2, 2, 2),
|
|
test_cfg=test_cfg,
|
|
**cfg)
|
|
|
|
model.requires_grad_(False)
|
|
return model
|
|
|
|
|
|
def test_ssd_head_forward():
|
|
"""Test SSD Head forward in torch and onnxruntime env."""
|
|
ssd_model = ssd_config()
|
|
|
|
featmap_size = [38, 19, 10, 6, 5, 3, 1]
|
|
|
|
feats = [
|
|
torch.rand(1, ssd_model.in_channels[i], featmap_size[i],
|
|
featmap_size[i]) for i in range(len(ssd_model.in_channels))
|
|
]
|
|
ort_validate(ssd_model.forward, feats)
|
|
|
|
|
|
def test_ssd_head_onnx_export():
|
|
"""Test SSD Head get_bboxes in torch and onnxruntime env."""
|
|
ssd_model = ssd_config()
|
|
s = 300
|
|
img_metas = [{
|
|
'img_shape_for_onnx': torch.Tensor([s, s]),
|
|
'scale_factor': np.ones(4),
|
|
'pad_shape': (s, s, 3),
|
|
'img_shape': (s, s, 2)
|
|
}]
|
|
|
|
# The data of ssd_head_get_bboxes.pkl contains two parts:
|
|
# cls_score(list(Tensor)) and bboxes(list(Tensor)),
|
|
# where each torch.Tensor is generated by torch.rand().
|
|
# the cls_score's size: (1, 20, 38, 38), (1, 30, 19, 19),
|
|
# (1, 30, 10, 10), (1, 30, 5, 5), (1, 20, 3, 3), (1, 20, 1, 1).
|
|
# the bboxes's size: (1, 16, 38, 38), (1, 24, 19, 19),
|
|
# (1, 24, 10, 10), (1, 24, 5, 5), (1, 16, 3, 3), (1, 16, 1, 1).
|
|
ssd_head_data = 'ssd_head_get_bboxes.pkl'
|
|
feats = mmcv.load(osp.join(data_path, ssd_head_data))
|
|
cls_score = feats[:6]
|
|
bboxes = feats[6:]
|
|
|
|
ssd_model.onnx_export = partial(
|
|
ssd_model.onnx_export, img_metas=img_metas, with_nms=False)
|
|
ort_validate(ssd_model.onnx_export, (cls_score, bboxes))
|