def _prepare_input_img(img_path, test_pipeline, shape=None): # build the data pipeline if shape is not None: test_pipeline[1]['img_scale'] = shape test_pipeline[1]['transforms'][0]['keep_ratio'] = False test_pipeline = [LoadImage()] + test_pipeline[1:] test_pipeline = Compose(test_pipeline) # prepare data data = dict(img=img_path) data = test_pipeline(data) imgs = data['img'] img_metas = [i.data for i in data['img_metas']] mm_inputs = {'imgs': imgs, 'img_metas': img_metas} return mm_inputs
def _prepare_input_img(img_path: str, test_pipeline: Iterable[dict], shape: Optional[Iterable] = None, rescale_shape: Optional[Iterable] = None) -> dict: # build the data pipeline if shape is not None: test_pipeline[1]['img_scale'] = (shape[1], shape[0]) test_pipeline[1]['transforms'][0]['keep_ratio'] = False test_pipeline = [LoadImage()] + test_pipeline[1:] test_pipeline = Compose(test_pipeline) # prepare data data = dict(img=img_path) data = test_pipeline(data) imgs = data['img'] img_metas = [i.data for i in data['img_metas']] if rescale_shape is not None: for img_meta in img_metas: img_meta['ori_shape'] = tuple(rescale_shape) + (3, ) mm_inputs = {'imgs': imgs, 'img_metas': img_metas} return mm_inputs
def pytorch2onnx(model, mm_inputs, opset_version=11, show=False, output_file='tmp.onnx', verify=False, dynamic_export=False): """Export Pytorch model to ONNX model and verify the outputs are same between Pytorch and ONNX. Args: model (nn.Module): Pytorch model we want to export. mm_inputs (dict): Contain the input tensors and img_metas information. opset_version (int): The onnx op version. Default: 11. show (bool): Whether print the computation graph. Default: False. output_file (string): The path to where we store the output ONNX model. Default: `tmp.onnx`. verify (bool): Whether compare the outputs between Pytorch and ONNX. Default: False. dynamic_export (bool): Whether to export ONNX with dynamic axis. Default: False. """ model.cpu().eval() test_mode = model.test_cfg.mode if isinstance(model.decode_head, nn.ModuleList): num_classes = model.decode_head[-1].num_classes else: num_classes = model.decode_head.num_classes imgs = mm_inputs.pop('imgs') img_metas = mm_inputs.pop('img_metas') img_list = [img[None, :] for img in imgs] img_meta_list = [[img_meta] for img_meta in img_metas] # update img_meta img_list, img_meta_list = _update_input_img(img_list, img_meta_list) # replace original forward function origin_forward = model.forward model.forward = partial( model.forward, img_metas=img_meta_list, return_loss=False, rescale=True) dynamic_axes = None if dynamic_export: if test_mode == 'slide': dynamic_axes = {'input': {0: 'batch'}, 'output': {1: 'batch'}} else: dynamic_axes = { 'input': { 0: 'batch', 2: 'height', 3: 'width' }, 'output': { 1: 'batch', 2: 'height', 3: 'width' } } register_extra_symbolics(opset_version) with torch.no_grad(): torch.onnx.export( model, (img_list, ), output_file, input_names=['input'], output_names=['output'], export_params=True, keep_initializers_as_inputs=False, verbose=show, opset_version=opset_version, dynamic_axes=dynamic_axes) print(f'Successfully exported ONNX model: {output_file}') model.forward = origin_forward if verify: # check by onnx import onnx onnx_model = onnx.load(output_file) onnx.checker.check_model(onnx_model) if dynamic_export and test_mode == 'whole': # scale image for dynamic shape test img_list = [resize(_, scale_factor=1.5) for _ in img_list] # concate flip image for batch test flip_img_list = [_.flip(-1) for _ in img_list] img_list = [ torch.cat((ori_img, flip_img), 0) for ori_img, flip_img in zip(img_list, flip_img_list) ] # update img_meta img_list, img_meta_list = _update_input_img( img_list, img_meta_list, test_mode == 'whole') # check the numerical value # get pytorch output with torch.no_grad(): pytorch_result = model(img_list, img_meta_list, return_loss=False) pytorch_result = np.stack(pytorch_result, 0) # get onnx output input_all = [node.name for node in onnx_model.graph.input] input_initializer = [ node.name for node in onnx_model.graph.initializer ] net_feed_input = list(set(input_all) - set(input_initializer)) assert (len(net_feed_input) == 1) sess = rt.InferenceSession(output_file) onnx_result = sess.run( None, {net_feed_input[0]: img_list[0].detach().numpy()})[0][0] # show segmentation results if show: import cv2 import os.path as osp img = img_meta_list[0][0]['filename'] if not osp.exists(img): img = imgs[0][:3, ...].permute(1, 2, 0) * 255 img = img.detach().numpy().astype(np.uint8) ori_shape = img.shape[:2] else: ori_shape = LoadImage()({'img': img})['ori_shape'] # resize onnx_result to ori_shape onnx_result_ = cv2.resize(onnx_result[0].astype(np.uint8), (ori_shape[1], ori_shape[0])) show_result_pyplot( model, img, (onnx_result_, ), palette=model.PALETTE, block=False, title='ONNXRuntime', opacity=0.5) # resize pytorch_result to ori_shape pytorch_result_ = cv2.resize(pytorch_result[0].astype(np.uint8), (ori_shape[1], ori_shape[0])) show_result_pyplot( model, img, (pytorch_result_, ), title='PyTorch', palette=model.PALETTE, opacity=0.5) # compare results np.testing.assert_allclose( pytorch_result.astype(np.float32) / num_classes, onnx_result.astype(np.float32) / num_classes, rtol=1e-5, atol=1e-5, err_msg='The outputs are different between Pytorch and ONNX') print('The outputs are same between Pytorch and ONNX')