Skip to content
New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

bug fix of multi batch size and support onnx of yolo-s model #296

Open
wants to merge 2 commits into
base: master
Choose a base branch
from

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
# for free to join this conversation on GitHub. Already have an account? # to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant