After converting to onnx, the dynamic batch reasoning results are inconsistent, but the static batch results are consistent

[sevenold]

mmcls 0.23.2
mmcv 1.7.0
mmcv-full 1.7.0
onnx 1.12.0
onnxruntime 1.9.0
torch 1.10.2

1. torch inference

import numpy as np
import torch

model = torch.load('model.pth')
model.eval()

dummy_input = torch.ones(2, 3, 480, 224)
_ = model(dummy_input)
# print(_)
for r in _:
    print(np.argmax(r), np.max(r), r[0].shape)
print('*'*100)
dummy_input = torch.ones(1, 3, 480, 224)
_ = model(dummy_input)
# print(_)
for r in _:
    print(np.argmax(r), np.max(r), r[0].shape)


2 0.99790525 ()
2 0.99790525 ()
****************************************************************************************************
2 0.99790525 ()

swin-large_16xb64_in1k.py

model = dict(
    type='ImageClassifier',
    backbone=dict(
        type='SwinTransformer',
        arch='large',
        img_size=(480, 224),
        drop_path_rate=0.5),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=11,
        in_channels=1536,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        topk=1),
    init_cfg=[
        dict(type='TruncNormal', layer='Linear', std=0.02, bias=0.0),
        dict(type='Constant', layer='LayerNorm', val=1.0, bias=0.0)
    ])
classes = [
    'normal0', 'normal1', 'normal2', 'porn0', 'gamble0', 'gamble1', 'gamble2',
    'fenfa0', 'fenfa1', 'dianshang0', 'dianshang1'
]
dataset_type = 'ImageNet'
img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='Resize',
        size=(480, 224),
        backend='pillow',
        interpolation='bicubic'),
    dict(
        type='Normalize',
        mean=[123.675, 116.28, 103.53],
        std=[58.395, 57.12, 57.375],
        to_rgb=True),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='Collect', keys=['img'])
]

2. save pth

from functools import partial

import mmcv
import torch
from mmcv.onnx import register_extra_symbolics
from mmcv.runner import load_checkpoint

from mmcls.models import build_classifier

torch.manual_seed(3)


if __name__ == '__main__':
    input_shape = (1, 3, 480, 224)
    cfg = mmcv.Config.fromfile('/image_classification/train_log/swin-large_16xb64_in1k.py')
    cfg.model.pretrained = None
    # build the model and load checkpoint
    classifier = build_classifier(cfg.model)
    model = load_checkpoint(classifier, '/image_classification/train_log/epoch_160.pth', map_location='cpu')

    model.cpu().eval()

    if hasattr(model.head, 'num_classes'):
        num_classes = model.head.num_classes
    # Some backbones use `num_classes=-1` to disable top classifier.
    elif getattr(model.backbone, 'num_classes', -1) > 0:
        num_classes = model.backbone.num_classes
    else:
        raise AttributeError('Cannot find "num_classes" in both head and backbone, please check the config file.')
    model.forward = partial(model.forward, img_metas={}, return_loss=False)
    register_extra_symbolics(13)

    torch.save(model, 'model.pth')

4. convert to onnx with dynamic_axes

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = torch.load("/image_classification/tools/deployment/model.pth", map_location=torch.device('cpu'))  # pytorch模型加载
batch_size = 8  # 批处理大小
input_shape = (3, 480, 224)  # 输入数据,改成自己的输入shape

# #set the model to inference mode
model.eval()

x = torch.ones(batch_size, *input_shape)  # 生成张量
x = x.to(device)
dynamic_axes = {
            'input': {
                0: 'batch',
            },
            'probs': {
                0: 'batch'
            }
        }
export_onnx_file = "test_1.onnx"  # 目的ONNX文件名
torch.onnx.export(
    model,
    x,
    export_onnx_file,
    input_names=['input'],
    output_names=['probs'],
    export_params=True,
    do_constant_folding=True,
    dynamic_axes=dynamic_axes, 
    opset_version=13)

5. test onnx with dynamic_axes

import numpy as np
import onnxruntime as ort
import onnx


model_dir = "test_1.onnx"
image = np.ones((3, 3, 480, 224)).astype(np.float32)
# predict
sess = ort.InferenceSession(model_dir)
res = sess.run([sess.get_outputs()[0].name], {sess.get_inputs()[0].name: image})
# decode
for r in res[0]:
    print(np.argmax(r), np.max(r), r[0].shape)
    # print(r)
print('*'*100)
image = np.ones((1, 3, 480, 224)).astype(np.float32)
# predict
res = sess.run([sess.get_outputs()[0].name], {sess.get_inputs()[0].name: image})
# decode
for r in res[0]:
    print(np.argmax(r), np.max(r), r[0].shape)

0 0.9803807 ()
0 0.742439 ()
0 0.51174736 ()
****************************************************************************************************
0 0.92829067 ()

6. convert to onnx without dynamic_axes

torch.onnx.export(
    model,
    x,
    export_onnx_file,
    input_names=['input'],
    output_names=['probs'],
    export_params=True,
    do_constant_folding=True,
  #  dynamic_axes=dynamic_axes, 
    opset_version=13)

7. test onnx without dynamic_axes

import numpy as np
import onnxruntime as ort
import onnx


model_dir = "test_1.onnx"
image = np.ones((8, 3, 480, 224)).astype(np.float32)
# predict
sess = ort.InferenceSession(model_dir)
res = sess.run([sess.get_outputs()[0].name], {sess.get_inputs()[0].name: image})
# decode
for r in res[0]:
    print(np.argmax(r), np.max(r), r[0].shape)
    # print(r)

2 0.99790525 ()
2 0.99790525 ()
2 0.99790525 ()
2 0.99790525 ()
2 0.99790525 ()
2 0.99790525 ()
2 0.99790525 ()
2 0.99790525 ()

Please help to find out where the problem is, thank you！！