mmdeploy/tests/test_codebase/test_mmdet3d/test_mmdet3d_models.py

198 lines
7.3 KiB
Python

# Copyright (c) OpenMMLab. All rights reserved.
import mmcv
import numpy as np
import pytest
import torch
from mmdeploy.codebase import import_codebase
from mmdeploy.utils import Backend, Codebase, Task, load_config
from mmdeploy.utils.test import WrapModel, check_backend, get_rewrite_outputs
try:
import_codebase(Codebase.MMDET3D)
except ImportError:
pytest.skip(
f'{Codebase.MMDET3D} is not installed.', allow_module_level=True)
model_cfg = load_config(
'tests/test_codebase/test_mmdet3d/data/model_cfg.py')[0]
def get_pillar_encoder():
from mmdet3d.models.voxel_encoders import PillarFeatureNet
model = PillarFeatureNet(
in_channels=4,
feat_channels=(64, ),
with_distance=False,
with_cluster_center=True,
with_voxel_center=True,
voxel_size=(0.2, 0.2, 4),
point_cloud_range=(0, -40, -3, 70.4, 40, 1),
norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
mode='max')
model.requires_grad_(False)
return model
def get_pointpillars_scatter():
from mmdet3d.models.middle_encoders import PointPillarsScatter
model = PointPillarsScatter(in_channels=64, output_shape=(16, 16))
model.requires_grad_(False)
return model
@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
def test_pillar_encoder(backend_type: Backend):
check_backend(backend_type, True)
model = get_pillar_encoder()
model.cpu().eval()
deploy_cfg = mmcv.Config(
dict(
backend_config=dict(type=backend_type.value),
onnx_config=dict(
input_shape=None,
input_names=['features', 'num_points', 'coors'],
output_names=['outputs']),
codebase_config=dict(
type=Codebase.MMDET3D.value, task=Task.VOXEL_DETECTION.value)))
features = torch.rand(3945, 32, 4) * 100
num_points = torch.randint(0, 32, (3945, ), dtype=torch.int32)
coors = torch.randint(0, 10, (3945, 4), dtype=torch.int32)
model_outputs = [model.forward(features, num_points, coors)]
wrapped_model = WrapModel(model, 'forward')
rewrite_inputs = {
'features': features,
'num_points': num_points,
'coors': coors
}
rewrite_outputs, is_backend_output = get_rewrite_outputs(
wrapped_model=wrapped_model,
model_inputs=rewrite_inputs,
deploy_cfg=deploy_cfg)
if isinstance(rewrite_outputs, dict):
rewrite_outputs = rewrite_outputs['output']
for model_output, rewrite_output in zip(model_outputs, rewrite_outputs):
if isinstance(rewrite_output, torch.Tensor):
rewrite_output = rewrite_output.cpu().numpy()
assert np.allclose(
model_output.shape, rewrite_output.shape, rtol=1e-03, atol=1e-03)
@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
def test_pointpillars_scatter(backend_type: Backend):
check_backend(backend_type, True)
model = get_pointpillars_scatter()
model.cpu().eval()
deploy_cfg = mmcv.Config(
dict(
backend_config=dict(type=backend_type.value),
onnx_config=dict(
input_shape=None,
input_names=['voxel_features', 'coors'],
output_names=['outputs']),
codebase_config=dict(
type=Codebase.MMDET3D.value, task=Task.VOXEL_DETECTION.value)))
voxel_features = torch.rand(16 * 16, 64) * 100
coors = torch.randint(0, 10, (16 * 16, 4), dtype=torch.int32)
model_outputs = [model.forward_batch(voxel_features, coors, 1)]
wrapped_model = WrapModel(model, 'forward_batch')
rewrite_inputs = {'voxel_features': voxel_features, 'coors': coors}
rewrite_outputs, is_backend_output = get_rewrite_outputs(
wrapped_model=wrapped_model,
model_inputs=rewrite_inputs,
deploy_cfg=deploy_cfg)
if isinstance(rewrite_outputs, dict):
rewrite_outputs = rewrite_outputs['output']
for model_output, rewrite_output in zip(model_outputs, rewrite_outputs):
if isinstance(rewrite_output, torch.Tensor):
rewrite_output = rewrite_output.cpu().numpy()
assert np.allclose(
model_output.shape, rewrite_output.shape, rtol=1e-03, atol=1e-03)
def get_centerpoint():
from mmdet3d.models.detectors.centerpoint import CenterPoint
model = CenterPoint(**model_cfg.centerpoint_model)
model.requires_grad_(False)
return model
def get_centerpoint_head():
from mmdet3d.models import builder
model_cfg.centerpoint_model.pts_bbox_head.test_cfg = model_cfg.\
centerpoint_model.test_cfg
head = builder.build_head(model_cfg.centerpoint_model.pts_bbox_head)
head.requires_grad_(False)
return head
@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
def test_centerpoint(backend_type: Backend):
from mmdeploy.codebase.mmdet3d.deploy.voxel_detection import VoxelDetection
from mmdeploy.core import RewriterContext
check_backend(backend_type, True)
model = get_centerpoint()
model.cpu().eval()
deploy_cfg = mmcv.Config(
dict(
backend_config=dict(type=backend_type.value),
onnx_config=dict(
input_shape=None,
opset_version=11,
input_names=['voxels', 'num_points', 'coors'],
output_names=['outputs']),
codebase_config=dict(
type=Codebase.MMDET3D.value, task=Task.VOXEL_DETECTION.value)))
voxeldetection = VoxelDetection(model_cfg, deploy_cfg, 'cpu')
inputs, data = voxeldetection.create_input(
'tests/test_codebase/test_mmdet3d/data/kitti/kitti_000008.bin')
with RewriterContext(
cfg=deploy_cfg,
backend=deploy_cfg.backend_config.type,
opset=deploy_cfg.onnx_config.opset_version):
outputs = model.forward(*data)
head = get_centerpoint_head()
rewrite_outputs = head.get_bboxes(*[[i] for i in outputs],
inputs['img_metas'][0])
assert rewrite_outputs is not None
def get_pointpillars_nus():
from mmdet3d.models.detectors import MVXFasterRCNN
model = MVXFasterRCNN(**model_cfg.pointpillars_nus_model)
model.requires_grad_(False)
return model
@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
def test_pointpillars_nus(backend_type: Backend):
from mmdeploy.codebase.mmdet3d.deploy.voxel_detection import VoxelDetection
from mmdeploy.core import RewriterContext
check_backend(backend_type, True)
model = get_pointpillars_nus()
model.cpu().eval()
deploy_cfg = mmcv.Config(
dict(
backend_config=dict(type=backend_type.value),
onnx_config=dict(
input_shape=None,
opset_version=11,
input_names=['voxels', 'num_points', 'coors'],
output_names=['outputs']),
codebase_config=dict(
type=Codebase.MMDET3D.value, task=Task.VOXEL_DETECTION.value)))
voxeldetection = VoxelDetection(model_cfg, deploy_cfg, 'cpu')
inputs, data = voxeldetection.create_input(
'tests/test_codebase/test_mmdet3d/data/kitti/kitti_000008.bin')
with RewriterContext(
cfg=deploy_cfg,
backend=deploy_cfg.backend_config.type,
opset=deploy_cfg.onnx_config.opset_version):
outputs = model.forward(*data)
assert outputs is not None