164 lines
4.7 KiB
Python
164 lines
4.7 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import os.path as osp
|
|
|
|
import mmcv
|
|
import pytest
|
|
import torch
|
|
|
|
from mmdet import digit_version
|
|
from mmdet.models.necks import FPN, YOLOV3Neck
|
|
from .utils import ort_validate
|
|
|
|
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)
|
|
|
|
# Control the returned model of fpn_neck_config()
|
|
fpn_test_step_names = {
|
|
'fpn_normal': 0,
|
|
'fpn_wo_extra_convs': 1,
|
|
'fpn_lateral_bns': 2,
|
|
'fpn_bilinear_upsample': 3,
|
|
'fpn_scale_factor': 4,
|
|
'fpn_extra_convs_inputs': 5,
|
|
'fpn_extra_convs_laterals': 6,
|
|
'fpn_extra_convs_outputs': 7,
|
|
}
|
|
|
|
# Control the returned model of yolo_neck_config()
|
|
yolo_test_step_names = {'yolo_normal': 0}
|
|
|
|
data_path = osp.join(osp.dirname(__file__), 'data')
|
|
|
|
|
|
def fpn_neck_config(test_step_name):
|
|
"""Return the class containing the corresponding attributes according to
|
|
the fpn_test_step_names."""
|
|
s = 64
|
|
in_channels = [8, 16, 32, 64]
|
|
feat_sizes = [s // 2**i for i in range(4)] # [64, 32, 16, 8]
|
|
out_channels = 8
|
|
|
|
feats = [
|
|
torch.rand(1, in_channels[i], feat_sizes[i], feat_sizes[i])
|
|
for i in range(len(in_channels))
|
|
]
|
|
|
|
if (fpn_test_step_names[test_step_name] == 0):
|
|
fpn_model = FPN(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_extra_convs=True,
|
|
num_outs=5)
|
|
elif (fpn_test_step_names[test_step_name] == 1):
|
|
fpn_model = FPN(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_extra_convs=False,
|
|
num_outs=5)
|
|
elif (fpn_test_step_names[test_step_name] == 2):
|
|
fpn_model = FPN(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_extra_convs=True,
|
|
no_norm_on_lateral=False,
|
|
norm_cfg=dict(type='BN', requires_grad=True),
|
|
num_outs=5)
|
|
elif (fpn_test_step_names[test_step_name] == 3):
|
|
fpn_model = FPN(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_extra_convs=True,
|
|
upsample_cfg=dict(mode='bilinear', align_corners=True),
|
|
num_outs=5)
|
|
elif (fpn_test_step_names[test_step_name] == 4):
|
|
fpn_model = FPN(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_extra_convs=True,
|
|
upsample_cfg=dict(scale_factor=2),
|
|
num_outs=5)
|
|
elif (fpn_test_step_names[test_step_name] == 5):
|
|
fpn_model = FPN(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_extra_convs='on_input',
|
|
num_outs=5)
|
|
elif (fpn_test_step_names[test_step_name] == 6):
|
|
fpn_model = FPN(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_extra_convs='on_lateral',
|
|
num_outs=5)
|
|
elif (fpn_test_step_names[test_step_name] == 7):
|
|
fpn_model = FPN(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_extra_convs='on_output',
|
|
num_outs=5)
|
|
return fpn_model, feats
|
|
|
|
|
|
def yolo_neck_config(test_step_name):
|
|
"""Config yolov3 Neck."""
|
|
in_channels = [16, 8, 4]
|
|
out_channels = [8, 4, 2]
|
|
|
|
# The data of yolov3_neck.pkl contains a list of
|
|
# torch.Tensor, where each torch.Tensor is generated by
|
|
# torch.rand and each tensor size is:
|
|
# (1, 4, 64, 64), (1, 8, 32, 32), (1, 16, 16, 16).
|
|
yolov3_neck_data = 'yolov3_neck.pkl'
|
|
feats = mmcv.load(osp.join(data_path, yolov3_neck_data))
|
|
|
|
if (yolo_test_step_names[test_step_name] == 0):
|
|
yolo_model = YOLOV3Neck(
|
|
in_channels=in_channels, out_channels=out_channels, num_scales=3)
|
|
return yolo_model, feats
|
|
|
|
|
|
def test_fpn_normal():
|
|
outs = fpn_neck_config('fpn_normal')
|
|
ort_validate(*outs)
|
|
|
|
|
|
def test_fpn_wo_extra_convs():
|
|
outs = fpn_neck_config('fpn_wo_extra_convs')
|
|
ort_validate(*outs)
|
|
|
|
|
|
def test_fpn_lateral_bns():
|
|
outs = fpn_neck_config('fpn_lateral_bns')
|
|
ort_validate(*outs)
|
|
|
|
|
|
def test_fpn_bilinear_upsample():
|
|
outs = fpn_neck_config('fpn_bilinear_upsample')
|
|
ort_validate(*outs)
|
|
|
|
|
|
def test_fpn_scale_factor():
|
|
outs = fpn_neck_config('fpn_scale_factor')
|
|
ort_validate(*outs)
|
|
|
|
|
|
def test_fpn_extra_convs_inputs():
|
|
outs = fpn_neck_config('fpn_extra_convs_inputs')
|
|
ort_validate(*outs)
|
|
|
|
|
|
def test_fpn_extra_convs_laterals():
|
|
outs = fpn_neck_config('fpn_extra_convs_laterals')
|
|
ort_validate(*outs)
|
|
|
|
|
|
def test_fpn_extra_convs_outputs():
|
|
outs = fpn_neck_config('fpn_extra_convs_outputs')
|
|
ort_validate(*outs)
|
|
|
|
|
|
def test_yolo_normal():
|
|
outs = yolo_neck_config('yolo_normal')
|
|
ort_validate(*outs)
|