Export VIT model to onnx

🤗Hub
1

I try to export VIT model to onnx with follow command
python -m transformers.onnx --model=google/vit-base-patch16-224 onnx/
But the output model have config

    "output": [{
        "name": "1617",
        "data_type": "TYPE_FP32",
        "dims": [2, 384],
        "label_filename": "",
        "is_shape_tensor": false
    }, {
        "name": "last_hidden_state",
        "data_type": "TYPE_FP32",
        "dims": [-1, -1, 384],
        "label_filename": "",
        "is_shape_tensor": false
    }],

why the output shape is 384? the input model have output shape (cls_classifier): Linear(in_features=384, out_features=1000, bias=True)?

2

cc @lewtun

3

Hi @phelan164 I’m not sure I can reproduce what you’re seeing - could you share how you’re obtaining the ONNX config?

Here’s what I get after exporting the model with your command:

import onnx

model = onnx.load("onnx/model.onnx")
model.graph.output
# [name: "last_hidden_state"
# type {
#   tensor_type {
#     elem_type: 1
#     shape {
#       dim {
#         dim_param: "batch"
#       }
#       dim {
#         dim_param: "sequence"
#       }
#       dim {
#         dim_value: 768
#       }
#     }
#   }
# }
# , name: "1601"
# type {
#   tensor_type {
#     elem_type: 1
#     shape {
#       dim {
#         dim_value: 2
#       }
#       dim {
#         dim_value: 768
#       }
#     }
#   }
# }
# ]

As you can see, the output dimension of the last hidden state is 768, which is what one expects for this model (no head, equivalent to VitModel)

1 Like