deepinsight / deepinsight/insightface
VIT Torch2Onnx Trick
- Dominant language
- Python
- Stars
- 29.7k
- Forks
- 6.1k
- PR merge metrics
- No merged PRs in 30d
Description
With lower version of Pytorch(1.6.0 for example), problem occurs when using exported ONNX model to inference with a dynamic input.
`2022-07-19 17:04:28.618850992 [E:onnxruntime:, sequential_executor.cc:333 Execute] Non-zero status code returned while running Reshape node. Name:'Reshape_5' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:43 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, std::vector&) gsl::narrow_cast(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{32,256,12,12}, requested shape:{1,256,144}`
It is a flattening operation in the VIT backbone, which is designed to flatten the last two dimensions of the tensor.
The model was exported with dummy input (1,3,112,112). To make dynamic batch size feasible, the input shape dim_param was changed to 'None' afterwards.
`graph.input[0].type.tensor_type.shape.dim[0].dim_param = 'None'` (see recognition/arcface_torch/torch2onnx.py)
However, the Reshape node does not convey this message. It hard codes the destined shape.

When input shape is dynamic (32, 256, 12, 12), the destined shape fails to become (32,256,144) and remains (1,256,144).
Pytorch 1.11.0 fixed this problem. For lower version users, using torch.reshape instead of flatten enables a correct conversion.
Original Code (recognition/arcface_torch/backbones/vit.py):
`x = self.proj(x).flatten(2).transpose(1, 2)`
Use Reshape:
`x = self.proj(x)`
`batch_size, channels, height, width = x.shape`
`x = torch.reshape(x, (batch_size, channels, height*width))`
`x = x.transpose(1,2)`
It passes the onnx_helper check and passes the MFR online test.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.