pt to onnx

[zg7759]

I found a problem when converting onnx, unlike the previous model of Adaface, the vector output of VIT is not normalized, and the recall of the measured effect after adding normalization is much higher than that of res101, but it also brings fewer false positives, I don't know if the conversion is correct, please have the code of the latest model pt2onnx.

[zg7759]

import torch

from cvlface.general_utils.config_utils import load_config from cvlface.research.recognition.code.run_v1.models import get_model

if name == "main":

inputs_shape = (1, 3, 112, 112)
x = torch.randn(inputs_shape)
y = torch.randn(1, 5, 2)
# setting 1: input is image
for config_name in [
    #'models/iresnet/configs/v1_ir50.yaml',
    #'D:/CVLface/cvlface/research/recognition/code/run_v1/models/vit/configs/v1_base.yaml',
    'D:/CVLface/cvlface/research/recognition/code/run_v1/models/vit_kprpe/configs/v1_base_kprpe_splithead_unshared.yaml',
    #     'models/swin_kprpe/configs/v1_base_kprpe_splithead_unshared.yaml',
    #'models/swin/configs/v1_base.yaml',
    #'models/vit_irpe/configs/v1_base_irpe.yaml',
    #'models/part_fvit/configs/v1_base.yaml',
]:
    config = load_config(config_name)
    config.yaml_path = config_name
    model = get_model(config, task='run_v1')
    #
    # model.set_swish(memory_efficient=False)
    model.load_state_dict_from_path('D:/CVLface/kvit.pt')
    # model.load_state_dict_from_path('vit.pt')
    model.eval()
# # * For onnx model
torch.onnx.export(
    model,
    (x, y),
    "D:/CVLface/kvit.onnx",
    input_names=["input", "input1"],
    output_names=["output"],
    do_constant_folding=True,
    keep_initializers_as_inputs=False,
    verbose=False,
    opset_version=13,
    export_params=True,
)      Here's my pt2onnx code, VIT's can convert inference, vit_kprpe inference results are wrong