Skip to content

Conversation

@wufei-png
Copy link

@wufei-png wufei-png commented May 4, 2024

fix: txt_feats should repeat batch_size times:
when export onnx model with batch size > 1 the img_feats shape is:

img_feats torch.Size([batch_size, 128, 80, 80])
img_feats torch.Size([batch_size, 256, 40, 40])
img_feats torch.Size([batch_size, 512, 20, 20])

however,the txt_feats is always : torch.Size([1, 80, 512]) it should be torch.Size([batch_size, 80, 512]) to avoid the export error like this:

File "/home/wufei2/anaconda3/envs/train10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wufei2/anaconda3/envs/train10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1488, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/wufei2/go/src/github.com/AILab-CVC/YOLO-World/deploy/../yolo_world/models/layers/yolo_bricks.py", line 289, in forward
    return x * self.scale + text_features
RuntimeError: The size of tensor a (40) must match the size of tensor b (80) at non-singleton dimension 1

chore: support onnx export:
I use YOLO-World-S model to export onnx format,have the err same as this issue:
I referenced the solution in this issue to add the max and avgpool that is compatible with onnx exports

Now export onnx has two compatible problems: einsum and pool layer, and I don't want to add a new bool variable (because other new compatibility problems may occur in future) so I replaced bool use_einsum with export_onnx.

The code has been tested in my local env.

-------------comment at 5.18:
@wondervictor
I notice that this pr have a small conflict about code formatting after recent main branch's code update, which has now been modified and is ready to be merged in:
image

@wufei-png wufei-png marked this pull request as draft May 4, 2024 06:23
@wufei-png wufei-png marked this pull request as ready for review May 4, 2024 06:42
@wufei-png wufei-png changed the title Wf/onnx bug fix of multi batch size and support onnx of yolo-s model May 4, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant