search

Tencent / TPAT

TensorRT Plugin Autogen Tool
Apache License 2.0
365 stars 42 forks source link

support for one hot plugin with dynamic axis other than batch_size dim #30

Open frankxyy opened 1 year ago

frankxyy commented 1 year ago

I read the code of one hot plugin generation. It seems that it only support dynamic axis of batch size dimension. Can tpat support one hot plugin generation with dynamic axises other than batch_size dimension?

wenqf11 commented 1 year ago

@frankxyy Currently we don't support dynamic batch_size on other dimensions. You may try to add Transpose before and after one_hot node in ONNX graph.

frankxyy commented 1 year ago

@wenqf11 Hi, why adding Transpose before and after one_hot node can avoid issues of dynamic dims? I suppose the transpose node cannot infer the exact tensort shape at runtime.

wenqf11 commented 1 year ago

@frankxyy If your ONNX model has more than one dynamic dimensions, we cannot support currently. Otherwise, only one dynamic diemnsion, you can Transpose it to the first dimension. It may cost more time to run it with Transpose.

frankxyy commented 1 year ago

@wenqf11 Thank you. Got it.