Reference for ultralytics/nn/modules/head.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/nn/modules/head.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!
ultralytics.nn.modules.head.Detect
Detect(nc: int = 80, ch: Tuple = ())
Bases: Module
YOLO Detect head for object detection models.
This class implements the detection head used in YOLO models for predicting bounding boxes and class probabilities. It supports both training and inference modes, with optional end-to-end detection capabilities.
Attributes:
Name | Type | Description |
---|---|---|
dynamic |
bool
|
Force grid reconstruction. |
export |
bool
|
Export mode flag. |
format |
str
|
Export format. |
end2end |
bool
|
End-to-end detection mode. |
max_det |
int
|
Maximum detections per image. |
shape |
tuple
|
Input shape. |
anchors |
Tensor
|
Anchor points. |
strides |
Tensor
|
Feature map strides. |
legacy |
bool
|
Backward compatibility for v3/v5/v8/v9 models. |
xyxy |
bool
|
Output format, xyxy or xywh. |
nc |
int
|
Number of classes. |
nl |
int
|
Number of detection layers. |
reg_max |
int
|
DFL channels. |
no |
int
|
Number of outputs per anchor. |
stride |
Tensor
|
Strides computed during build. |
cv2 |
ModuleList
|
Convolution layers for box regression. |
cv3 |
ModuleList
|
Convolution layers for classification. |
dfl |
Module
|
Distribution Focal Loss layer. |
one2one_cv2 |
ModuleList
|
One-to-one convolution layers for box regression. |
one2one_cv3 |
ModuleList
|
One-to-one convolution layers for classification. |
Methods:
Name | Description |
---|---|
forward |
Perform forward pass and return predictions. |
forward_end2end |
Perform forward pass for end-to-end detection. |
bias_init |
Initialize detection head biases. |
decode_bboxes |
Decode bounding boxes from predictions. |
postprocess |
Post-process model predictions. |
Examples:
Create a detection head for 80 classes
>>> detect = Detect(nc=80, ch=(256, 512, 1024))
>>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
>>> outputs = detect(x)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
nc
|
int
|
Number of classes. |
80
|
ch
|
tuple
|
Tuple of channel sizes from backbone feature maps. |
()
|
Source code in ultralytics/nn/modules/head.py
78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
|
bias_init
bias_init()
Initialize Detect() biases, WARNING: requires stride availability.
Source code in ultralytics/nn/modules/head.py
191 192 193 194 195 196 197 198 199 200 201 202 |
|
decode_bboxes
decode_bboxes(
bboxes: Tensor, anchors: Tensor, xywh: bool = True
) -> torch.Tensor
Decode bounding boxes from predictions.
Source code in ultralytics/nn/modules/head.py
204 205 206 |
|
forward
forward(x: List[Tensor]) -> Union[List[torch.Tensor], Tuple]
Concatenate and return predicted bounding boxes and class probabilities.
Source code in ultralytics/nn/modules/head.py
114 115 116 117 118 119 120 121 122 123 124 |
|
forward_end2end
forward_end2end(x: List[Tensor]) -> Union[dict, Tuple]
Perform forward pass of the v10Detect module.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
x
|
List[Tensor]
|
Input feature maps from different levels. |
required |
Returns:
Name | Type | Description |
---|---|---|
outputs |
dict | tuple
|
Training mode returns dict with one2many and one2one outputs. Inference mode returns processed detections or tuple with detections and raw outputs. |
Source code in ultralytics/nn/modules/head.py
126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
|
postprocess
staticmethod
postprocess(preds: Tensor, max_det: int, nc: int = 80) -> torch.Tensor
Post-process YOLO model predictions.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
preds
|
Tensor
|
Raw predictions with shape (batch_size, num_anchors, 4 + nc) with last dimension format [x, y, w, h, class_probs]. |
required |
max_det
|
int
|
Maximum detections per image. |
required |
nc
|
int
|
Number of classes. |
80
|
Returns:
Type | Description |
---|---|
Tensor
|
Processed predictions with shape (batch_size, min(max_det, num_anchors), 6) and last dimension format [x, y, w, h, max_class_prob, class_index]. |
Source code in ultralytics/nn/modules/head.py
208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 |
|
ultralytics.nn.modules.head.Segment
Segment(nc: int = 80, nm: int = 32, npr: int = 256, ch: Tuple = ())
Bases: Detect
YOLO Segment head for segmentation models.
This class extends the Detect head to include mask prediction capabilities for instance segmentation tasks.
Attributes:
Name | Type | Description |
---|---|---|
nm |
int
|
Number of masks. |
npr |
int
|
Number of protos. |
proto |
Proto
|
Prototype generation module. |
cv4 |
ModuleList
|
Convolution layers for mask coefficients. |
Methods:
Name | Description |
---|---|
forward |
Return model outputs and mask coefficients. |
Examples:
Create a segmentation head
>>> segment = Segment(nc=80, nm=32, npr=256, ch=(256, 512, 1024))
>>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
>>> outputs = segment(x)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
nc
|
int
|
Number of classes. |
80
|
nm
|
int
|
Number of masks. |
32
|
npr
|
int
|
Number of protos. |
256
|
ch
|
tuple
|
Tuple of channel sizes from backbone feature maps. |
()
|
Source code in ultralytics/nn/modules/head.py
255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 |
|
forward
forward(x: List[Tensor]) -> Union[Tuple, List[torch.Tensor]]
Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients.
Source code in ultralytics/nn/modules/head.py
273 274 275 276 277 278 279 280 281 282 |
|
ultralytics.nn.modules.head.OBB
OBB(nc: int = 80, ne: int = 1, ch: Tuple = ())
Bases: Detect
YOLO OBB detection head for detection with rotation models.
This class extends the Detect head to include oriented bounding box prediction with rotation angles.
Attributes:
Name | Type | Description |
---|---|---|
ne |
int
|
Number of extra parameters. |
cv4 |
ModuleList
|
Convolution layers for angle prediction. |
angle |
Tensor
|
Predicted rotation angles. |
Methods:
Name | Description |
---|---|
forward |
Concatenate and return predicted bounding boxes and class probabilities. |
decode_bboxes |
Decode rotated bounding boxes. |
Examples:
Create an OBB detection head
>>> obb = OBB(nc=80, ne=1, ch=(256, 512, 1024))
>>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
>>> outputs = obb(x)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
nc
|
int
|
Number of classes. |
80
|
ne
|
int
|
Number of extra parameters. |
1
|
ch
|
tuple
|
Tuple of channel sizes from backbone feature maps. |
()
|
Source code in ultralytics/nn/modules/head.py
307 308 309 310 311 312 313 314 315 316 317 318 319 320 |
|
decode_bboxes
decode_bboxes(bboxes: Tensor, anchors: Tensor) -> torch.Tensor
Decode rotated bounding boxes.
Source code in ultralytics/nn/modules/head.py
336 337 338 |
|
forward
forward(x: List[Tensor]) -> Union[torch.Tensor, Tuple]
Concatenate and return predicted bounding boxes and class probabilities.
Source code in ultralytics/nn/modules/head.py
322 323 324 325 326 327 328 329 330 331 332 333 334 |
|
ultralytics.nn.modules.head.Pose
Pose(nc: int = 80, kpt_shape: Tuple = (17, 3), ch: Tuple = ())
Bases: Detect
YOLO Pose head for keypoints models.
This class extends the Detect head to include keypoint prediction capabilities for pose estimation tasks.
Attributes:
Name | Type | Description |
---|---|---|
kpt_shape |
tuple
|
Number of keypoints and dimensions (2 for x,y or 3 for x,y,visible). |
nk |
int
|
Total number of keypoint values. |
cv4 |
ModuleList
|
Convolution layers for keypoint prediction. |
Methods:
Name | Description |
---|---|
forward |
Perform forward pass through YOLO model and return predictions. |
kpts_decode |
Decode keypoints from predictions. |
Examples:
Create a pose detection head
>>> pose = Pose(nc=80, kpt_shape=(17, 3), ch=(256, 512, 1024))
>>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
>>> outputs = pose(x)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
nc
|
int
|
Number of classes. |
80
|
kpt_shape
|
tuple
|
Number of keypoints, number of dims (2 for x,y or 3 for x,y,visible). |
(17, 3)
|
ch
|
tuple
|
Tuple of channel sizes from backbone feature maps. |
()
|
Source code in ultralytics/nn/modules/head.py
363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 |
|
forward
forward(x: List[Tensor]) -> Union[torch.Tensor, Tuple]
Perform forward pass through YOLO model and return predictions.
Source code in ultralytics/nn/modules/head.py
379 380 381 382 383 384 385 386 387 |
|
kpts_decode
kpts_decode(bs: int, kpts: Tensor) -> torch.Tensor
Decode keypoints from predictions.
Source code in ultralytics/nn/modules/head.py
389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 |
|
ultralytics.nn.modules.head.Classify
Classify(
c1: int,
c2: int,
k: int = 1,
s: int = 1,
p: Optional[int] = None,
g: int = 1,
)
Bases: Module
YOLO classification head, i.e. x(b,c1,20,20) to x(b,c2).
This class implements a classification head that transforms feature maps into class predictions.
Attributes:
Name | Type | Description |
---|---|---|
export |
bool
|
Export mode flag. |
conv |
Conv
|
Convolutional layer for feature transformation. |
pool |
AdaptiveAvgPool2d
|
Global average pooling layer. |
drop |
Dropout
|
Dropout layer for regularization. |
linear |
Linear
|
Linear layer for final classification. |
Methods:
Name | Description |
---|---|
forward |
Perform forward pass of the YOLO model on input image data. |
Examples:
Create a classification head
>>> classify = Classify(c1=1024, c2=1000)
>>> x = torch.randn(1, 1024, 20, 20)
>>> output = classify(x)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
c1
|
int
|
Number of input channels. |
required |
c2
|
int
|
Number of output classes. |
required |
k
|
int
|
Kernel size. |
1
|
s
|
int
|
Stride. |
1
|
p
|
int
|
Padding. |
None
|
g
|
int
|
Groups. |
1
|
Source code in ultralytics/nn/modules/head.py
444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 |
|
forward
forward(x: Union[List[Tensor], Tensor]) -> Union[torch.Tensor, Tuple]
Perform forward pass of the YOLO model on input image data.
Source code in ultralytics/nn/modules/head.py
463 464 465 466 467 468 469 470 471 |
|
ultralytics.nn.modules.head.WorldDetect
WorldDetect(
nc: int = 80, embed: int = 512, with_bn: bool = False, ch: Tuple = ()
)
Bases: Detect
Head for integrating YOLO detection models with semantic understanding from text embeddings.
This class extends the standard Detect head to incorporate text embeddings for enhanced semantic understanding in object detection tasks.
Attributes:
Name | Type | Description |
---|---|---|
cv3 |
ModuleList
|
Convolution layers for embedding features. |
cv4 |
ModuleList
|
Contrastive head layers for text-vision alignment. |
Methods:
Name | Description |
---|---|
forward |
Concatenate and return predicted bounding boxes and class probabilities. |
bias_init |
Initialize detection head biases. |
Examples:
Create a WorldDetect head
>>> world_detect = WorldDetect(nc=80, embed=512, with_bn=False, ch=(256, 512, 1024))
>>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
>>> text = torch.randn(1, 80, 512)
>>> outputs = world_detect(x, text)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
nc
|
int
|
Number of classes. |
80
|
embed
|
int
|
Embedding dimension. |
512
|
with_bn
|
bool
|
Whether to use batch normalization in contrastive head. |
False
|
ch
|
tuple
|
Tuple of channel sizes from backbone feature maps. |
()
|
Source code in ultralytics/nn/modules/head.py
497 498 499 500 501 502 503 504 505 506 507 508 509 510 |
|
bias_init
bias_init()
Initialize Detect() biases, WARNING: requires stride availability.
Source code in ultralytics/nn/modules/head.py
522 523 524 525 526 527 528 |
|
forward
forward(x: List[Tensor], text: Tensor) -> Union[List[torch.Tensor], Tuple]
Concatenate and return predicted bounding boxes and class probabilities.
Source code in ultralytics/nn/modules/head.py
512 513 514 515 516 517 518 519 520 |
|
ultralytics.nn.modules.head.LRPCHead
LRPCHead(vocab: Module, pf: Module, loc: Module, enabled: bool = True)
Bases: Module
Lightweight Region Proposal and Classification Head for efficient object detection.
This head combines region proposal filtering with classification to enable efficient detection with dynamic vocabulary support.
Attributes:
Name | Type | Description |
---|---|---|
vocab |
Module
|
Vocabulary/classification layer. |
pf |
Module
|
Proposal filter module. |
loc |
Module
|
Localization module. |
enabled |
bool
|
Whether the head is enabled. |
Methods:
Name | Description |
---|---|
conv2linear |
Convert a 1x1 convolutional layer to a linear layer. |
forward |
Process classification and localization features to generate detection proposals. |
Examples:
Create an LRPC head
>>> vocab = nn.Conv2d(256, 80, 1)
>>> pf = nn.Conv2d(256, 1, 1)
>>> loc = nn.Conv2d(256, 4, 1)
>>> head = LRPCHead(vocab, pf, loc, enabled=True)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
vocab
|
Module
|
Vocabulary/classification module. |
required |
pf
|
Module
|
Proposal filter module. |
required |
loc
|
Module
|
Localization module. |
required |
enabled
|
bool
|
Whether to enable the head functionality. |
True
|
Source code in ultralytics/nn/modules/head.py
557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 |
|
conv2linear
conv2linear(conv: Conv2d) -> nn.Linear
Convert a 1x1 convolutional layer to a linear layer.
Source code in ultralytics/nn/modules/head.py
573 574 575 576 577 578 579 |
|
forward
forward(
cls_feat: Tensor, loc_feat: Tensor, conf: float
) -> Tuple[Tuple, torch.Tensor]
Process classification and localization features to generate detection proposals.
Source code in ultralytics/nn/modules/head.py
581 582 583 584 585 586 587 588 589 590 591 592 593 594 |
|
ultralytics.nn.modules.head.YOLOEDetect
YOLOEDetect(
nc: int = 80, embed: int = 512, with_bn: bool = False, ch: Tuple = ()
)
Bases: Detect
Head for integrating YOLO detection models with semantic understanding from text embeddings.
This class extends the standard Detect head to support text-guided detection with enhanced semantic understanding through text embeddings and visual prompt embeddings.
Attributes:
Name | Type | Description |
---|---|---|
is_fused |
bool
|
Whether the model is fused for inference. |
cv3 |
ModuleList
|
Convolution layers for embedding features. |
cv4 |
ModuleList
|
Contrastive head layers for text-vision alignment. |
reprta |
Residual
|
Residual block for text prompt embeddings. |
savpe |
SAVPE
|
Spatial-aware visual prompt embeddings module. |
embed |
int
|
Embedding dimension. |
Methods:
Name | Description |
---|---|
fuse |
Fuse text features with model weights for efficient inference. |
get_tpe |
Get text prompt embeddings with normalization. |
get_vpe |
Get visual prompt embeddings with spatial awareness. |
forward_lrpc |
Process features with fused text embeddings for prompt-free model. |
forward |
Process features with class prompt embeddings to generate detections. |
bias_init |
Initialize biases for detection heads. |
Examples:
Create a YOLOEDetect head
>>> yoloe_detect = YOLOEDetect(nc=80, embed=512, with_bn=True, ch=(256, 512, 1024))
>>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
>>> cls_pe = torch.randn(1, 80, 512)
>>> outputs = yoloe_detect(x, cls_pe)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
nc
|
int
|
Number of classes. |
80
|
embed
|
int
|
Embedding dimension. |
512
|
with_bn
|
bool
|
Whether to use batch normalization in contrastive head. |
False
|
ch
|
tuple
|
Tuple of channel sizes from backbone feature maps. |
()
|
Source code in ultralytics/nn/modules/head.py
630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 |
|
bias_init
bias_init()
Initialize biases for detection heads.
Source code in ultralytics/nn/modules/head.py
775 776 777 778 779 780 781 782 783 784 |
|
forward
forward(
x: List[Tensor], cls_pe: Tensor, return_mask: bool = False
) -> Union[torch.Tensor, Tuple]
Process features with class prompt embeddings to generate detections.
Source code in ultralytics/nn/modules/head.py
761 762 763 764 765 766 767 768 769 770 771 772 773 |
|
forward_lrpc
forward_lrpc(
x: List[Tensor], return_mask: bool = False
) -> Union[torch.Tensor, Tuple]
Process features with fused text embeddings to generate detections for prompt-free model.
Source code in ultralytics/nn/modules/head.py
723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 |
|
fuse
fuse(txt_feats: Tensor)
Fuse text features with model weights for efficient inference.
Source code in ultralytics/nn/modules/head.py
663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 |
|
get_tpe
get_tpe(tpe: Optional[Tensor]) -> Optional[torch.Tensor]
Get text prompt embeddings with normalization.
Source code in ultralytics/nn/modules/head.py
710 711 712 |
|
get_vpe
get_vpe(x: List[Tensor], vpe: Tensor) -> torch.Tensor
Get visual prompt embeddings with spatial awareness.
Source code in ultralytics/nn/modules/head.py
714 715 716 717 718 719 720 721 |
|
ultralytics.nn.modules.head.YOLOESegment
YOLOESegment(
nc: int = 80,
nm: int = 32,
npr: int = 256,
embed: int = 512,
with_bn: bool = False,
ch: Tuple = (),
)
Bases: YOLOEDetect
YOLO segmentation head with text embedding capabilities.
This class extends YOLOEDetect to include mask prediction capabilities for instance segmentation tasks with text-guided semantic understanding.
Attributes:
Name | Type | Description |
---|---|---|
nm |
int
|
Number of masks. |
npr |
int
|
Number of protos. |
proto |
Proto
|
Prototype generation module. |
cv5 |
ModuleList
|
Convolution layers for mask coefficients. |
Methods:
Name | Description |
---|---|
forward |
Return model outputs and mask coefficients. |
Examples:
Create a YOLOESegment head
>>> yoloe_segment = YOLOESegment(nc=80, nm=32, npr=256, embed=512, with_bn=True, ch=(256, 512, 1024))
>>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
>>> text = torch.randn(1, 80, 512)
>>> outputs = yoloe_segment(x, text)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
nc
|
int
|
Number of classes. |
80
|
nm
|
int
|
Number of masks. |
32
|
npr
|
int
|
Number of protos. |
256
|
embed
|
int
|
Embedding dimension. |
512
|
with_bn
|
bool
|
Whether to use batch normalization in contrastive head. |
False
|
ch
|
tuple
|
Tuple of channel sizes from backbone feature maps. |
()
|
Source code in ultralytics/nn/modules/head.py
811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 |
|
forward
forward(x: List[Tensor], text: Tensor) -> Union[Tuple, torch.Tensor]
Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients.
Source code in ultralytics/nn/modules/head.py
833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 |
|
ultralytics.nn.modules.head.RTDETRDecoder
RTDETRDecoder(
nc: int = 80,
ch: Tuple = (512, 1024, 2048),
hd: int = 256,
nq: int = 300,
ndp: int = 4,
nh: int = 8,
ndl: int = 6,
d_ffn: int = 1024,
dropout: float = 0.0,
act: Module = nn.ReLU(),
eval_idx: int = -1,
nd: int = 100,
label_noise_ratio: float = 0.5,
box_noise_scale: float = 1.0,
learnt_init_query: bool = False,
)
Bases: Module
Real-Time Deformable Transformer Decoder (RTDETRDecoder) module for object detection.
This decoder module utilizes Transformer architecture along with deformable convolutions to predict bounding boxes and class labels for objects in an image. It integrates features from multiple layers and runs through a series of Transformer decoder layers to output the final predictions.
Attributes:
Name | Type | Description |
---|---|---|
export |
bool
|
Export mode flag. |
hidden_dim |
int
|
Dimension of hidden layers. |
nhead |
int
|
Number of heads in multi-head attention. |
nl |
int
|
Number of feature levels. |
nc |
int
|
Number of classes. |
num_queries |
int
|
Number of query points. |
num_decoder_layers |
int
|
Number of decoder layers. |
input_proj |
ModuleList
|
Input projection layers for backbone features. |
decoder |
DeformableTransformerDecoder
|
Transformer decoder module. |
denoising_class_embed |
Embedding
|
Class embeddings for denoising. |
num_denoising |
int
|
Number of denoising queries. |
label_noise_ratio |
float
|
Label noise ratio for training. |
box_noise_scale |
float
|
Box noise scale for training. |
learnt_init_query |
bool
|
Whether to learn initial query embeddings. |
tgt_embed |
Embedding
|
Target embeddings for queries. |
query_pos_head |
MLP
|
Query position head. |
enc_output |
Sequential
|
Encoder output layers. |
enc_score_head |
Linear
|
Encoder score prediction head. |
enc_bbox_head |
MLP
|
Encoder bbox prediction head. |
dec_score_head |
ModuleList
|
Decoder score prediction heads. |
dec_bbox_head |
ModuleList
|
Decoder bbox prediction heads. |
Methods:
Name | Description |
---|---|
forward |
Run forward pass and return bounding box and classification scores. |
Examples:
Create an RTDETRDecoder
>>> decoder = RTDETRDecoder(nc=80, ch=(512, 1024, 2048), hd=256, nq=300)
>>> x = [torch.randn(1, 512, 64, 64), torch.randn(1, 1024, 32, 32), torch.randn(1, 2048, 16, 16)]
>>> outputs = decoder(x)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
nc
|
int
|
Number of classes. |
80
|
ch
|
tuple
|
Channels in the backbone feature maps. |
(512, 1024, 2048)
|
hd
|
int
|
Dimension of hidden layers. |
256
|
nq
|
int
|
Number of query points. |
300
|
ndp
|
int
|
Number of decoder points. |
4
|
nh
|
int
|
Number of heads in multi-head attention. |
8
|
ndl
|
int
|
Number of decoder layers. |
6
|
d_ffn
|
int
|
Dimension of the feed-forward networks. |
1024
|
dropout
|
float
|
Dropout rate. |
0.0
|
act
|
Module
|
Activation function. |
ReLU()
|
eval_idx
|
int
|
Evaluation index. |
-1
|
nd
|
int
|
Number of denoising. |
100
|
label_noise_ratio
|
float
|
Label noise ratio. |
0.5
|
box_noise_scale
|
float
|
Box noise scale. |
1.0
|
learnt_init_query
|
bool
|
Whether to learn initial query embeddings. |
False
|
Source code in ultralytics/nn/modules/head.py
898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 |
|
forward
forward(
x: List[Tensor], batch: Optional[dict] = None
) -> Union[Tuple, torch.Tensor]
Run the forward pass of the module, returning bounding box and classification scores for the input.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
x
|
List[Tensor]
|
List of feature maps from the backbone. |
required |
batch
|
dict
|
Batch information for training. |
None
|
Returns:
Name | Type | Description |
---|---|---|
outputs |
tuple | Tensor
|
During training, returns a tuple of bounding boxes, scores, and other metadata. During inference, returns a tensor of shape (bs, 300, 4+nc) containing bounding boxes and class scores. |
Source code in ultralytics/nn/modules/head.py
977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 |
|
ultralytics.nn.modules.head.v10Detect
v10Detect(nc: int = 80, ch: Tuple = ())
Bases: Detect
v10 Detection head from https://arxiv.org/pdf/2405.14458.
This class implements the YOLOv10 detection head with dual-assignment training and consistent dual predictions for improved efficiency and performance.
Attributes:
Name | Type | Description |
---|---|---|
end2end |
bool
|
End-to-end detection mode. |
max_det |
int
|
Maximum number of detections. |
cv3 |
ModuleList
|
Light classification head layers. |
one2one_cv3 |
ModuleList
|
One-to-one classification head layers. |
Methods:
Name | Description |
---|---|
forward |
Perform forward pass of the v10Detect module. |
bias_init |
Initialize biases of the Detect module. |
fuse |
Remove the one2many head for inference optimization. |
Examples:
Create a v10Detect head
>>> v10_detect = v10Detect(nc=80, ch=(256, 512, 1024))
>>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
>>> outputs = v10_detect(x)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
nc
|
int
|
Number of classes. |
80
|
ch
|
tuple
|
Tuple of channel sizes from backbone feature maps. |
()
|
Source code in ultralytics/nn/modules/head.py
1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 |
|
fuse
fuse()
Remove the one2many head for inference optimization.
Source code in ultralytics/nn/modules/head.py
1226 1227 1228 |
|