import torch import torch.nn as nn import torch.nn.functional as F try: from mmcv.tensorrt import (TRTWrapper, is_tensorrt_plugin_loaded, onnx2trt, save_trt_engine) except ImportError: pytest.skip('TensorRT should be installed from source.', allow_module_level=True) if not torch.cuda.is_available(): pytest.skip('CUDA is required for this test module', allow_module_level=True) if not is_tensorrt_plugin_loaded(): pytest.skip('Test requires to complie TensorRT plugins in mmcv', allow_module_level=True) class WrapFunction(nn.Module): def __init__(self, wrapped_function): super(WrapFunction, self).__init__() self.wrapped_function = wrapped_function def forward(self, *args, **kwargs): return self.wrapped_function(*args, **kwargs) onnx_file = 'tmp.onnx' trt_file = 'tmp.engine'
help='Dataset name') parser.add_argument('--verify', action='store_true', help='Verify the outputs of ONNXRuntime and TensorRT') parser.add_argument( '--verbose', action='store_true', help='Whether to verbose logging messages while creating \ TensorRT engine.') args = parser.parse_args() return args if __name__ == '__main__': assert is_tensorrt_plugin_loaded(), 'TensorRT plugin should be compiled.' args = parse_args() if not args.input_img: args.input_img = osp.join(osp.dirname(__file__), '../demo/demo.png') # check arguments assert osp.exists(args.config), 'Config {} not found.'.format(args.config) assert osp.exists(args.model), \ 'ONNX model {} not found.'.format(args.model) assert args.workspace_size >= 0, 'Workspace size less than 0.' assert DATASETS.get(args.dataset) is not None, \ 'Dataset {} does not found.'.format(args.dataset) for max_value, min_value in zip(args.max_shape, args.min_shape): assert max_value >= min_value, \ 'max_shape should be larger than min shape'
def test_roialign(): try: from mmcv.tensorrt import (TRTWraper, onnx2trt, save_trt_engine, is_tensorrt_plugin_loaded) if not is_tensorrt_plugin_loaded(): pytest.skip('test requires to complie TensorRT plugins in mmcv') except (ImportError, ModuleNotFoundError): pytest.skip('test requires to install TensorRT from source.') try: from mmcv.ops import RoIAlign except (ImportError, ModuleNotFoundError): pytest.skip('test requires compilation') # trt config fp16_mode = False max_workspace_size = 1 << 30 # roi align config pool_h = 2 pool_w = 2 spatial_scale = 1.0 sampling_ratio = 2 inputs = [([[[[1., 2.], [3., 4.]]]], [[0., 0., 0., 1., 1.]]), ([[[[1., 2.], [3., 4.]], [[4., 3.], [2., 1.]]]], [[0., 0., 0., 1., 1.]]), ([[[[1., 2., 5., 6.], [3., 4., 7., 8.], [9., 10., 13., 14.], [11., 12., 15., 16.]]]], [[0., 0., 0., 3., 3.]])] wrapped_model = RoIAlign((pool_w, pool_h), spatial_scale, sampling_ratio, 'avg', True).cuda() for case in inputs: np_input = np.array(case[0], dtype=np.float32) np_rois = np.array(case[1], dtype=np.float32) input = torch.from_numpy(np_input).cuda() rois = torch.from_numpy(np_rois).cuda() with torch.no_grad(): torch.onnx.export( wrapped_model, (input, rois), onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['input', 'rois'], output_names=['roi_feat'], opset_version=11) onnx_model = onnx.load(onnx_file) # create trt engine and wraper opt_shape_dict = { 'input': [list(input.shape), list(input.shape), list(input.shape)], 'rois': [list(rois.shape), list(rois.shape), list(rois.shape)] } trt_engine = onnx2trt( onnx_model, opt_shape_dict, fp16_mode=fp16_mode, max_workspace_size=max_workspace_size) save_trt_engine(trt_engine, trt_file) trt_model = TRTWraper(trt_file, ['input', 'rois'], ['roi_feat']) with torch.no_grad(): trt_outputs = trt_model({'input': input, 'rois': rois}) trt_roi_feat = trt_outputs['roi_feat'] # compute pytorch_output with torch.no_grad(): pytorch_roi_feat = wrapped_model(input, rois) # allclose if os.path.exists(onnx_file): os.remove(onnx_file) if os.path.exists(trt_file): os.remove(trt_file) assert torch.allclose(pytorch_roi_feat, trt_roi_feat)