Prune YOLOWorld head with C2fAttn Modules?

[edoardocastagnini]

Hi everyone,

I’m working on YOLOWorld, and I’ve successfully applied pruning across the backbone using Torch-Pruning’s tools. However, I’m struggling with pruning the model’s HEAD, specifically modules of type C2fAttn, which include the MaxSigmoidAttnBlock.
Torch-Pruning’s ViT-style head pruning methods are incompatible and often causes shape mismatches, so I was wondering if anyone already explored pruning strategies for YOLOWorld’s head, or is there guidance on this scenario?

[UltralyticsAssistant]

👋 Hello @edoardocastagnini, thank you for your interest in Ultralytics 🚀! This is an automated response — an Ultralytics engineer will also assist you here shortly.

We recommend a visit to the Docs for general guidance where you can find many Python and CLI usage examples and where many of the most common questions may already be answered:

• Docs: https://docs.ultralytics.com/
• Python: https://docs.ultralytics.com/usage/python/
• CLI: https://docs.ultralytics.com/usage/cli/

If this is a 🐛 Bug Report, please provide a minimum reproducible example (MRE) to help us debug it: https://docs.ultralytics.com/help/minimum-reproducible-example/

If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples and training logs, and verify you are following our Tips for Best Training Results: https://docs.ultralytics.com/guides/model-training-tips/

Join the Ultralytics community where it suits you best. For real-time chat, head to Discord 🎧: https://discord.com/invite/ultralytics. Prefer in-depth discussions? Check out Discourse: https://community.ultralytics.com/. Or dive into threads on our Subreddit to share knowledge with the community: https://www.reddit.com/r/Ultralytics/

Upgrade

Upgrade to the latest ultralytics package including all requirements in a Python>=3.8 environment with PyTorch>=1.8 to verify your issue is not already resolved in the latest version:

pip install -U ultralytics

• Requirements: https://github.com/ultralytics/ultralytics/blob/main/pyproject.toml

Environments

YOLO may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide: https://docs.ultralytics.com/yolov5/environments/google_cloud_quickstart_tutorial/
• Amazon Deep Learning AMI. See AWS Quickstart Guide: https://docs.ultralytics.com/yolov5/environments/aws_quickstart_tutorial/
• Docker Image. See Docker Quickstart Guide: https://docs.ultralytics.com/yolov5/environments/docker_image_quickstart_tutorial/

Status

If this badge is green, all Ultralytics CI tests are currently passing. CI tests verify correct operation of all YOLO Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

To help us assist you faster regarding pruning the YOLOWorld head and C2fAttn/MaxSigmoidAttnBlock specifics:

• Please share a minimal reproducible example that triggers the shape mismatch
• Include your ultralytics and PyTorch versions, model YAML, and pruning script snippet
• Add any error logs/tracebacks and the exact layers you’re targeting

Thanks again for the detailed context — we’ll take a closer look soon! 🙌

[glenn-jocher]

Great question — Torch-Pruning’s ViT rules won’t map cleanly because C2fAttn’s MaxSigmoidAttnBlock reshapes channels into nh heads, so pruning must preserve these invariants: c2 % nh == 0, ec == c2, and prune proj_conv.out_channels, the embedding conv out_channels (ec), and the guide Linear(gc→ec) together in groups of hc = c2 // nh channels per head; in practice either prune whole heads (drop 1×hc groups across those layers and decrement nh) or exclude C2fAttn from pruning and focus on adjacent convs instead. See the MaxSigmoidAttnBlock reference and the C2fAttn reference for expected shapes and coupling.