deepinsight / deepinsight/insightface

VIT Torch2Onnx Trick

Open
#2,052 4 comments 0 reactions 0 assignees View on GitHub
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.
![image](https://user-images.githubusercontent.com/62649623/179912603-9fde9409-7c0e-44a2-b7c2-d44406af3625.png)
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.