Failed to convert a dynamic input model containing torch.nn.TransformerDecoderLayer

[lzcchl]

can you give me some advice?

lib version:
torch 2.3.0
onnx 1.17.0
tensorrt 10.8

this is code I get onnx:

import torch
import torch.nn as nn
from torch import Tensor
from torch.nn import TransformerDecoderLayer

class MyModel(nn.Module):
    def __init__(self, embed_dim: int=512, num_heads: int=8) -> None:
        super().__init__()
        self.layers = TransformerDecoderLayer(embed_dim, num_heads, batch_first=True)
        return 
    
    def forward(self, img_features: Tensor, txt_features: Tensor) -> Tensor:
        attn_output = self.layers(img_features, txt_features)
        return attn_output
    
if __name__ == '__main__':
    output_onnx_path = "./test.onnx"

    img_features = torch.randn((1,8,512))
    txt_features = torch.randn((1,16,512))

    model = MyModel()
    model.eval()

    out = model(img_features, txt_features)
    print('out.shape: {}'.format(out.shape))

    input_names = ["img_features", "txt_features"]
    dynamic = {"img_features": {0: "batch", 1: "number"},
            "txt_features": {0: "batch", 1: "seq_len"},}
    output_names = ['output']
    dynamic['output'] = {0: "batch", 1: "number"}

    torch.onnx.export(
        model.cpu(),
        (img_features.cpu(), txt_features.cpu()),
        output_onnx_path,
        verbose=False,
        opset_version=17,
        do_constant_folding=True,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic,
    )
    
    print('finish')

and this is bash I convert to trt:

trt_exe=/usr/src/tensorrt/bin/trtexec

onnx_path="./test.onnx"
trt_path="./test.trt"

$trt_exe --onnx=$onnx_path \
    --saveEngine=$trt_path \
    --skipInference \
    --minShapes="img_features:1x8x512,txt_features:1x16x512" \
    --optShapes="img_features:1x16x512,txt_features:1x32x512" \
    --maxShapes="img_features:1x16x512,txt_features:1x32x512" \

I get Error like this:

Error[4]: IBuilder::buildSerializedNetwork: Error Code 4: Internal Error (kOPT values for profile 0 violate shape constraints: IShuffleLayer /layers/self_attn/Reshape_4: reshaping failed for tensor: /layers/self_attn/Gather_4_output_0 reshape would change volume 8192 to 4096)

[lzcchl]

I solve it, maybe torch.nn.TransformerDecoderLayer use torch.randn((8,512)) rather than torch.randn((1,8,512))