Skip to content

Fix visual prompt training for YOLOE26#23401

Merged
Laughing-q merged 5 commits into
ultralytics:mainfrom
ShuaiLYU:fix/visual-prompt-training-for-yoloe26
Jan 24, 2026
Merged

Fix visual prompt training for YOLOE26#23401
Laughing-q merged 5 commits into
ultralytics:mainfrom
ShuaiLYU:fix/visual-prompt-training-for-yoloe26

Conversation

@ShuaiLYU

@ShuaiLYU ShuaiLYU commented Jan 23, 2026

Copy link
Copy Markdown
Contributor

Fix Visual Prompt Training and Typo in YOLOE26

Summary

This PR fixes visual prompt training issues YOLOE/YOLOE26 models.

Changes

1. Add tal_topk2 parameter support to TVPDetectLoss

  • Added tal_topk2 parameter to TVPDetectLoss.__init__() for flexible task-aligned assignment configuration
  • Passes tal_topk2 through to the underlying v8DetectionLoss for better one-to-one assignment control
  • Removed unnecessary score slicing in _get_vp_features to use the full score tensor

2. Fix typo in YOLOEDetect forward_lrpc

  • Corrected self.one2one_cv2 to self.one2one_cv3 in the forward_lrpc method

🛠️ PR Summary

Made with ❤️ by Ultralytics Actions

🌟 Summary

Fixes visual-prompt training behavior for YOLOe26 by correcting an end-to-end head routing bug and aligning TVP loss feature/parameter handling. 🛠️

📊 Key Changes

  • Corrected forward_lrpc() to use one2one_cv3 (not one2one_cv2) for classification features in end-to-end mode.
  • Extended TVPDetectLoss to accept and forward an optional tal_topk2 parameter into v8DetectionLoss.
  • Updated VP feature extraction to use full preds["scores"] instead of slicing by ori_nc, preventing unintended channel dropping.

🎯 Purpose & Impact

  • Prevents miswired heads during end-to-end training, improving correctness and stability for prompt-based detection. ✅
  • Ensures TVP/visual-prompt loss uses the intended score tensor layout, avoiding shape/class mismatches. 🎯
  • Improves configurability of task-aligned assignment via tal_topk2, helping advanced training setups behave as expected. ⚙️

@codecov

codecov Bot commented Jan 23, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 25.00000% with 3 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
ultralytics/utils/loss.py 33.33% 2 Missing ⚠️
ultralytics/nn/modules/head.py 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@ShuaiLYU

ShuaiLYU commented Jan 23, 2026

Copy link
Copy Markdown
Contributor Author

Test demo @Laughing-q

from ultralytics import YOLO
from ultralytics.models.yolo.yoloe import YOLOEVPTrainer

model=YOLO("yoloe-26s-seg.yaml").load("yoloe-26s-seg.pt")

data=dict(train=dict(yolo_data=["coco128.yaml"]),
          
          val=dict(yolo_data=["coco128.yaml"]))


model.train(
    data=data,
    task='detect',
    epochs=10,
    batch=4,
    imgsz=640,
    workers=4,
    single_cls=False,
    trainer=YOLOEVPTrainer)

@UltralyticsAssistant UltralyticsAssistant added bug Something isn't working as intended in the official Ultralytics package. detect Object Detection issues, PR's fixed Bug has been resolved labels Jan 23, 2026
@UltralyticsAssistant

Copy link
Copy Markdown
Member

👋 Hello @ShuaiLYU, thank you for submitting a ultralytics/ultralytics 🚀 PR! To ensure a seamless integration of your work, please review the following checklist:

-✅ Define a Purpose: Clearly explain the purpose of your fix or feature in your PR description, and link to any relevant issues. Ensure your commit messages are clear, concise, and adhere to the project's conventions.
-✅ Synchronize with Source: Confirm your PR is synchronized with the ultralytics/ultralytics main branch. If it's behind, update it by clicking the 'Update branch' button or by running git pull and git merge main locally.
-✅ Ensure CI Checks Pass: Verify all Ultralytics Continuous Integration (CI) checks are passing. If any checks fail, please address the issues.
-✅ Update Documentation: Update the relevant documentation for any new or modified features.
-✅ Add Tests: If applicable, include or update tests to cover your changes, and confirm that all tests are passing.
-✅ Sign the CLA: Please ensure you have signed our Contributor License Agreement if this is your first Ultralytics PR by writing "I have read the CLA Document and I sign the CLA" in a new message.
-✅ Minimize Changes: Limit your changes to the minimum necessary for your bug fix or feature addition. "It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is." — Bruce Lee

For more guidance, please refer to our Contributing Guide. This is an automated message, and an Ultralytics engineer will assist shortly. 🚀

@UltralyticsAssistant UltralyticsAssistant left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔍 PR Review

Made with ❤️ by Ultralytics Actions

Looks clean. The changes are straightforward and align with the intended behavior; no issues found.

@Laughing-q Laughing-q left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ShuaiLYU Nice fix!

@Laughing-q Laughing-q changed the title Fix/visual prompt training for yoloe26 Fix visual prompt training for YOLOE26 Jan 24, 2026
@Laughing-q Laughing-q merged commit 597e615 into ultralytics:main Jan 24, 2026
19 checks passed
kuazhangxiaoai pushed a commit to kuazhangxiaoai/ultralytics-semantic-segment that referenced this pull request Apr 10, 2026
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Jing Qiu <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working as intended in the official Ultralytics package. detect Object Detection issues, PR's fixed Bug has been resolved

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants