ultralytics 8.2.38 official YOLOv10 support (#13113)
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Laughing-q <1185102784@qq.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
parent
821e5fa477
commit
ffb46fd7fb
23 changed files with 785 additions and 32 deletions
|
|
@ -148,7 +148,7 @@ class KeypointLoss(nn.Module):
|
|||
class v8DetectionLoss:
|
||||
"""Criterion class for computing training losses."""
|
||||
|
||||
def __init__(self, model): # model must be de-paralleled
|
||||
def __init__(self, model, tal_topk=10): # model must be de-paralleled
|
||||
"""Initializes v8DetectionLoss with the model, defining model-related properties and BCE loss function."""
|
||||
device = next(model.parameters()).device # get model device
|
||||
h = model.args # hyperparameters
|
||||
|
|
@ -164,7 +164,7 @@ class v8DetectionLoss:
|
|||
|
||||
self.use_dfl = m.reg_max > 1
|
||||
|
||||
self.assigner = TaskAlignedAssigner(topk=10, num_classes=self.nc, alpha=0.5, beta=6.0)
|
||||
self.assigner = TaskAlignedAssigner(topk=tal_topk, num_classes=self.nc, alpha=0.5, beta=6.0)
|
||||
self.bbox_loss = BboxLoss(m.reg_max - 1, use_dfl=self.use_dfl).to(device)
|
||||
self.proj = torch.arange(m.reg_max, dtype=torch.float, device=device)
|
||||
|
||||
|
|
@ -714,3 +714,21 @@ class v8OBBLoss(v8DetectionLoss):
|
|||
b, a, c = pred_dist.shape # batch, anchors, channels
|
||||
pred_dist = pred_dist.view(b, a, 4, c // 4).softmax(3).matmul(self.proj.type(pred_dist.dtype))
|
||||
return torch.cat((dist2rbox(pred_dist, pred_angle, anchor_points), pred_angle), dim=-1)
|
||||
|
||||
|
||||
class E2EDetectLoss:
|
||||
"""Criterion class for computing training losses."""
|
||||
|
||||
def __init__(self, model):
|
||||
"""Initialize E2EDetectLoss with one-to-many and one-to-one detection losses using the provided model."""
|
||||
self.one2many = v8DetectionLoss(model, tal_topk=10)
|
||||
self.one2one = v8DetectionLoss(model, tal_topk=1)
|
||||
|
||||
def __call__(self, preds, batch):
|
||||
"""Calculate the sum of the loss for box, cls and dfl multiplied by batch size."""
|
||||
preds = preds[1] if isinstance(preds, tuple) else preds
|
||||
one2many = preds["one2many"]
|
||||
loss_one2many = self.one2many(one2many, batch)
|
||||
one2one = preds["one2one"]
|
||||
loss_one2one = self.one2one(one2one, batch)
|
||||
return loss_one2many[0] + loss_one2one[0], loss_one2many[1] + loss_one2one[1]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue