Repository: open-mmlab/mmpose

Issue:
About pytorch2onnx
Hello! I used your resnet18 demo and trained a customized dataset. It works and the trained model can recognize my images. 
Then, I use pytorch2onnx.py to transfer my models into .onnx format. The script worked and the ".onnx" model can generate the same result with the ".pth" one. It seems work.

![resnetonnx](https://user-images.githubusercontent.com/39186017/100060472-8acafc80-2e67-11eb-8ce7-8749e4f78bda.png)


But when I use netron the visualize the model structure, the resnet18 model seems strange: some modules contain some unexpected streams, which is obviously not the structure of resnet18. 
![resnetnetron](https://user-images.githubusercontent.com/39186017/100060166-2871fc00-2e67-11eb-8fd6-83b07296b8ab.png)

I wonder whether it is normal in onnx conversion? Whether it will be wrong in the following conversions like "ncnn"  or "tnn"?

-----------------------------------Edit---------------------------------------
Similarly, when I am converting mobilenet, some unexpected modules appear, though the result is correct. And I also have no idea about the correctness about the model. 
![mobilenetron](https://user-images.githubusercontent.com/39186017/100067241-fbc2e200-2e70-11eb-9682-a51adabd4e88.png)




Discussion:

Resolving patch:
diff --git a/tools/pytorch2onnx.py b/tools/pytorch2onnx.py
new file mode 100644
--- /dev/null
+++ b/tools/pytorch2onnx.py
@@ -0,0 +1,151 @@
+import argparse
+
+import mmcv
+import numpy as np
+import torch
+from mmcv.runner import load_checkpoint
+
+from mmpose.models import build_posenet
+
+try:
+    import onnx
+    import onnxruntime as rt
+except ImportError as e:
+    raise ImportError(f'Please install onnx and onnxruntime first. {e}')
+
+try:
+    from mmcv.onnx.symbolic import register_extra_symbolics
+except ModuleNotFoundError:
+    raise NotImplementedError('please update mmcv to version>=1.0.4')
+
+
+def _convert_batchnorm(module):
+    """Convert the syncBNs into normal BN3ds."""
+    module_output = module
+    if isinstance(module, torch.nn.SyncBatchNorm):
+        module_output = torch.nn.BatchNorm3d(module.num_features, module.eps,
+                                             module.momentum, module.affine,
+                                             module.track_running_stats)
+        if module.affine:
+            module_output.weight.data = module.weight.data.clone().detach()
+            module_output.bias.data = module.bias.data.clone().detach()
+            # keep requires_grad unchanged
+            module_output.weight.requires_grad = module.weight.requires_grad
+            module_output.bias.requires_grad = module.bias.requires_grad
+        module_output.running_mean = module.running_mean
+        module_output.running_var = module.running_var
+        module_output.num_batches_tracked = module.num_batches_tracked
+    for name, child in module.named_children():
+        module_output.add_module(name, _convert_batchnorm(child))
+    del module
+    return module_output
+
+
+def pytorch2onnx(model,
+                 input_shape,
+                 opset_version=11,
+                 show=False,
+                 output_file='tmp.onnx',
+                 verify=False):
+    """Convert pytorch model to onnx model.
+
+    Args:
+        model (:obj:`nn.Module`): The pytorch model to be exported.
+        input_shape (tuple[int]): The input tensor shape of the model.
+        opset_version (int): Opset version of onnx used. Default: 11.
+        show (bool): Determines whether to print the onnx model architecture.
+            Default: False.
+        output_file (str): Output onnx model name. Default: 'tmp.onnx'.
+        verify (bool): Determines whether to verify the onnx model.
+            Default: False.
+    """
+    model.cpu().eval()
+
+    one_img = torch.randn(input_shape)
+
+    register_extra_symbolics(opset_version)
+    torch.onnx.export(
+        model,
+        one_img,
+        output_file,
+        export_params=True,
+        keep_initializers_as_inputs=True,
+        verbose=show,
+        opset_version=opset_version)
+
+    print(f'Successfully exported ONNX model: {output_file}')
+    if verify:
+        # check by onnx
+        onnx_model = onnx.load(output_file)
+        onnx.checker.check_model(onnx_model)
+
+        # check the numerical value
+        # get pytorch output
+        pytorch_result = model(one_img).detach().numpy()
+
+        # 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]: one_img.detach().numpy()})[0]
+        # only compare part of results
+        assert np.allclose(
+            pytorch_result, onnx_result,
+            atol=1.e-5), 'The outputs are different between Pytorch and ONNX'
+        print('The numerical values are same between Pytorch and ONNX')
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='Convert MMPose models to ONNX')
+    parser.add_argument('config', help='test config file path')
+    parser.add_argument('checkpoint', help='checkpoint file')
+    parser.add_argument('--show', action='store_true', help='show onnx graph')
+    parser.add_argument('--output-file', type=str, default='tmp.onnx')
+    parser.add_argument('--opset-version', type=int, default=11)
+    parser.add_argument(
+        '--verify',
+        action='store_true',
+        help='verify the onnx model output against pytorch output')
+    parser.add_argument(
+        '--shape',
+        type=int,
+        nargs='+',
+        default=[1, 3, 256, 192],
+        help='input size')
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    assert args.opset_version == 11, 'MMPose only supports opset 11 now'
+
+    cfg = mmcv.Config.fromfile(args.config)
+    # build the model
+    model = build_posenet(cfg.model)
+    model = _convert_batchnorm(model)
+
+    # onnx.export does not support kwargs
+    if hasattr(model, 'forward_dummy'):
+        model.forward = model.forward_dummy
+    else:
+        raise NotImplementedError(
+            'Please implement the forward method for exporting.')
+
+    checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
+
+    # conver model to onnx file
+    pytorch2onnx(
+        model,
+        args.shape,
+        opset_version=args.opset_version,
+        show=args.show,
+        output_file=args.output_file,
+        verify=args.verify)


Test patch:
diff --git a/tests/test_onnx.py b/tests/test_onnx.py
new file mode 100644
--- /dev/null
+++ b/tests/test_onnx.py
@@ -0,0 +1,28 @@
+import os.path as osp
+import tempfile
+
+import torch.nn as nn
+from tools.pytorch2onnx import _convert_batchnorm, pytorch2onnx
+
+
+class TestModel(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv3d(1, 2, 1)
+        self.bn = nn.SyncBatchNorm(2)
+
+    def forward(self, x):
+        return self.bn(self.conv(x))
+
+    def forward_dummy(self, x):
+        return (self.forward(x), )
+
+
+def test_onnx_exporting():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        out_file = osp.join(tmpdir, 'tmp.onnx')
+        model = TestModel()
+        model = _convert_batchnorm(model)
+        # test exporting
+        pytorch2onnx(model, (1, 1, 1, 1, 1), output_file=out_file)