vendor/ViTPose/tests/test_onnx.py

# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import tempfile

import torch.nn as nn

from tools.deployment.pytorch2onnx import _convert_batchnorm, pytorch2onnx


class DummyModel(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 = DummyModel()
        model = _convert_batchnorm(model)
        # test exporting
        pytorch2onnx(model, (1, 1, 1, 1, 1), output_file=out_file)