From a953207f60f0f5bdf167f370e3b8b1a4cfd79ba9 Mon Sep 17 00:00:00 2001 From: Laughing-q <1185102784@qq.com> Date: Thu, 19 Sep 2024 17:36:39 +0800 Subject: [PATCH] revert kpt loss --- ultralytics/utils/loss.py | 107 +++++++++----------------------------- 1 file changed, 24 insertions(+), 83 deletions(-) diff --git a/ultralytics/utils/loss.py b/ultralytics/utils/loss.py index 9342756d7d..38cedd34f3 100644 --- a/ultralytics/utils/loss.py +++ b/ultralytics/utils/loss.py @@ -137,11 +137,11 @@ class KeypointLoss(nn.Module): def forward(self, pred_kpts, gt_kpts, kpt_mask, area): """Calculates keypoint loss factor and Euclidean distance loss for predicted and actual keypoints.""" - d = (pred_kpts[..., 0] - gt_kpts[..., 0]).pow(2) + (pred_kpts[..., 1] - gt_kpts[..., 1]).pow(2) - kpt_loss_factor = kpt_mask.shape[1] / (torch.sum(kpt_mask != 0, dim=1) + 1e-9) + d = (pred_kpts[..., 0] - gt_kpts[..., 0]) ** 2 + (pred_kpts[..., 1] - gt_kpts[..., 1]) ** 2 + kpt_loss_factor = (torch.sum(kpt_mask != 0) + torch.sum(kpt_mask == 0)) / (torch.sum(kpt_mask != 0) + 1e-9) # e = d / (2 * (area * self.sigmas) ** 2 + 1e-9) # from formula - e = d / ((2 * self.sigmas).pow(2) * (area + 1e-9) * 2) # from cocoeval - return (kpt_loss_factor.view(-1, 1) * ((1 - torch.exp(-e)) * kpt_mask)).mean() + e = d / (2 * self.sigmas) ** 2 / (area + 1e-9) / 2 # from cocoeval + return kpt_loss_factor * ((1 - torch.exp(-e)) * kpt_mask).mean() class v8DetectionLoss: @@ -437,7 +437,6 @@ class v8PoseLoss(v8DetectionLoss): """Criterion class for computing training losses.""" def __init__(self, model): # model must be de-paralleled - """Initializes v8PoseLoss with model, sets keypoint variables and declares a keypoint loss instance.""" super().__init__(model) self.kpt_shape = model.model[-1].kpt_shape self.bce_pose = nn.BCEWithLogitsLoss() @@ -454,7 +453,7 @@ class v8PoseLoss(v8DetectionLoss): (self.reg_max * 4, self.nc), 1 ) - # B, grids, .. + # b, grids, .. pred_scores = pred_scores.permute(0, 2, 1).contiguous() pred_distri = pred_distri.permute(0, 2, 1).contiguous() pred_kpts = pred_kpts.permute(0, 2, 1).contiguous() @@ -463,7 +462,7 @@ class v8PoseLoss(v8DetectionLoss): imgsz = torch.tensor(feats[0].shape[2:], device=self.device, dtype=dtype) * self.stride[0] # image size (h,w) anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5) - # Targets + # targets batch_size = pred_scores.shape[0] batch_idx = batch["batch_idx"].view(-1, 1) targets = torch.cat((batch_idx, batch["cls"].view(-1, 1), batch["bboxes"]), 1) @@ -471,7 +470,7 @@ class v8PoseLoss(v8DetectionLoss): gt_labels, gt_bboxes = targets.split((1, 4), 2) # cls, xyxy mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0) - # Pboxes + # pboxes pred_bboxes = self.bbox_decode(anchor_points, pred_distri) # xyxy, (b, h*w, 4) pred_kpts = self.kpts_decode(anchor_points, pred_kpts.view(batch_size, -1, *self.kpt_shape)) # (b, h*w, 17, 3) @@ -486,11 +485,11 @@ class v8PoseLoss(v8DetectionLoss): target_scores_sum = max(target_scores.sum(), 1) - # Cls loss + # cls loss # loss[1] = self.varifocal_loss(pred_scores, target_scores, target_labels) / target_scores_sum # VFL way loss[3] = self.bce(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum # BCE - # Bbox loss + # bbox loss if fg_mask.sum(): target_bboxes /= stride_tensor loss[0], loss[4] = self.bbox_loss( @@ -499,14 +498,23 @@ class v8PoseLoss(v8DetectionLoss): keypoints = batch["keypoints"].to(self.device).float().clone() keypoints[..., 0] *= imgsz[1] keypoints[..., 1] *= imgsz[0] - - loss[1], loss[2] = self.calculate_keypoints_loss( - fg_mask, target_gt_idx, keypoints, batch_idx, stride_tensor, target_bboxes, pred_kpts - ) + for i in range(batch_size): + if fg_mask[i].sum(): + idx = target_gt_idx[i][fg_mask[i]] + gt_kpt = keypoints[batch_idx.view(-1) == i][idx] # (n, 51) + gt_kpt[..., 0] /= stride_tensor[fg_mask[i]] + gt_kpt[..., 1] /= stride_tensor[fg_mask[i]] + area = xyxy2xywh(target_bboxes[i][fg_mask[i]])[:, 2:].prod(1, keepdim=True) + pred_kpt = pred_kpts[i][fg_mask[i]] + kpt_mask = gt_kpt[..., 2] != 0 + loss[1] += self.keypoint_loss(pred_kpt, gt_kpt, kpt_mask, area) # pose loss + # kpt_score loss + if pred_kpt.shape[-1] == 3: + loss[2] += self.bce_pose(pred_kpt[..., 2], kpt_mask.float()) # keypoint obj loss loss[0] *= self.hyp.box # box gain - loss[1] *= self.hyp.pose # pose gain - loss[2] *= self.hyp.kobj # kobj gain + loss[1] *= self.hyp.pose / batch_size # pose gain + loss[2] *= self.hyp.kobj / batch_size # kobj gain loss[3] *= self.hyp.cls # cls gain loss[4] *= self.hyp.dfl # dfl gain @@ -521,73 +529,6 @@ class v8PoseLoss(v8DetectionLoss): y[..., 1] += anchor_points[:, [1]] - 0.5 return y - def calculate_keypoints_loss( - self, masks, target_gt_idx, keypoints, batch_idx, stride_tensor, target_bboxes, pred_kpts - ): - """ - Calculate the keypoints loss for the model. - - This function calculates the keypoints loss and keypoints object loss for a given batch. The keypoints loss is - based on the difference between the predicted keypoints and ground truth keypoints. The keypoints object loss is - a binary classification loss that classifies whether a keypoint is present or not. - - Args: - masks (torch.Tensor): Binary mask tensor indicating object presence, shape (BS, N_anchors). - target_gt_idx (torch.Tensor): Index tensor mapping anchors to ground truth objects, shape (BS, N_anchors). - keypoints (torch.Tensor): Ground truth keypoints, shape (N_kpts_in_batch, N_kpts_per_object, kpts_dim). - batch_idx (torch.Tensor): Batch index tensor for keypoints, shape (N_kpts_in_batch, 1). - stride_tensor (torch.Tensor): Stride tensor for anchors, shape (N_anchors, 1). - target_bboxes (torch.Tensor): Ground truth boxes in (x1, y1, x2, y2) format, shape (BS, N_anchors, 4). - pred_kpts (torch.Tensor): Predicted keypoints, shape (BS, N_anchors, N_kpts_per_object, kpts_dim). - - Returns: - (tuple): Returns a tuple containing: - - kpts_loss (torch.Tensor): The keypoints loss. - - kpts_obj_loss (torch.Tensor): The keypoints object loss. - """ - batch_idx = batch_idx.flatten() - batch_size = len(masks) - - # Find the maximum number of keypoints in a single image - max_kpts = torch.unique(batch_idx, return_counts=True)[1].max() - - # Create a tensor to hold batched keypoints - batched_keypoints = torch.zeros( - (batch_size, max_kpts, keypoints.shape[1], keypoints.shape[2]), device=keypoints.device - ) - - # TODO: any idea how to vectorize this? - # Fill batched_keypoints with keypoints based on batch_idx - for i in range(batch_size): - keypoints_i = keypoints[batch_idx == i] - batched_keypoints[i, : keypoints_i.shape[0]] = keypoints_i - - # Expand dimensions of target_gt_idx to match the shape of batched_keypoints - target_gt_idx_expanded = target_gt_idx.unsqueeze(-1).unsqueeze(-1) - - # Use target_gt_idx_expanded to select keypoints from batched_keypoints - selected_keypoints = batched_keypoints.gather( - 1, target_gt_idx_expanded.expand(-1, -1, keypoints.shape[1], keypoints.shape[2]) - ) - - # Divide coordinates by stride - selected_keypoints /= stride_tensor.view(1, -1, 1, 1) - - kpts_loss = 0 - kpts_obj_loss = 0 - - if masks.any(): - gt_kpt = selected_keypoints[masks] - area = xyxy2xywh(target_bboxes[masks])[:, 2:].prod(1, keepdim=True) - pred_kpt = pred_kpts[masks] - kpt_mask = gt_kpt[..., 2] != 0 if gt_kpt.shape[-1] == 3 else torch.full_like(gt_kpt[..., 0], True) - kpts_loss = self.keypoint_loss(pred_kpt, gt_kpt, kpt_mask, area) # pose loss - - if pred_kpt.shape[-1] == 3: - kpts_obj_loss = self.bce_pose(pred_kpt[..., 2], kpt_mask.float()) # keypoint obj loss - - return kpts_loss, kpts_obj_loss - class v8ClassificationLoss: """Criterion class for computing training losses."""