Implementando YOLO desde cero en PyTorch

				
					import torch
import torch.nn as nn
from .box_ops import xywh_to_xyxy, ciou_loss_xyxy

class YOLOLoss(nn.Module):
    def __init__(self, anchors, strides, num_classes, img_size,
                 lambda_box=7.5, lambda_obj=1.0, lambda_cls=1.0,
                 obj_pos_weight=1.0, obj_neg_weight=0.5):
        super().__init__()
        self.anchors = anchors
        self.strides = strides
        self.num_classes = num_classes
        self.img_size = img_size

        self.lambda_box = lambda_box
        self.lambda_obj = lambda_obj
        self.lambda_cls = lambda_cls

        self.bce = nn.BCEWithLogitsLoss(reduction="none")
        self.obj_pos_weight = obj_pos_weight
        self.obj_neg_weight = obj_neg_weight

    def decode_scale(self, pred, scale_idx):
        """
        pred: [B, A*(5+C), S, S]
        returns:
          boxes_xyxy: [B, A, S, S, 4] in pixels
          obj_logit:  [B, A, S, S]
          cls_logit:  [B, A, S, S, C]
        """
        B, _, S, _ = pred.shape
        A = len(self.anchors[scale_idx])
        C = self.num_classes
        stride = self.strides[scale_idx]

        pred = pred.view(B, A, 5 + C, S, S).permute(0, 1, 3, 4, 2).contiguous()
        # [B, A, S, S, 5+C]
        tx_ty = pred[..., 0:2]
        tw_th = pred[..., 2:4]
        obj = pred[..., 4]
        cls = pred[..., 5:]

        # grid
        gy, gx = torch.meshgrid(torch.arange(S, device=pred.device),
                                torch.arange(S, device=pred.device), indexing="ij")
        grid = torch.stack([gx, gy], dim=-1).float()  # [S,S,2]

        # anchors
        anc = torch.tensor(self.anchors[scale_idx], device=pred.device).float()  # [A,2]
        anc = anc.view(1, A, 1, 1, 2)

        # decode center
        pxy = (tx_ty.sigmoid() + grid.view(1, 1, S, S, 2)) * stride  # pixels
        # decode wh
        pwh = (tw_th.exp() * anc)  # pixels

        boxes_xywh = torch.cat([pxy, pwh], dim=-1)
        boxes_xyxy = xywh_to_xyxy(boxes_xywh)
        return boxes_xyxy, obj, cls

    def forward(self, preds, targets_per_scale):
        """
        preds: list of 3 scale outputs
        targets_per_scale: list of (tbox, tobj, tcls) per scale
        """
        total_box = torch.tensor(0.0, device=preds[0].device)
        total_obj = torch.tensor(0.0, device=preds[0].device)
        total_cls = torch.tensor(0.0, device=preds[0].device)

        for s, pred in enumerate(preds):
            tbox, tobj, tcls = targets_per_scale[s]
            pbox, pobj_logit, pcls_logit = self.decode_scale(pred, s)

            # objectness loss with weighting
            obj_loss = self.bce(pobj_logit, tobj)
            w = torch.where(tobj > 0.5,
                            torch.full_like(obj_loss, self.obj_pos_weight),
                            torch.full_like(obj_loss, self.obj_neg_weight))
            total_obj = total_obj + (obj_loss * w).mean()

            # positives mask
            pos = tobj > 0.5
            if pos.any():
                # box loss CIoU
                pbox_pos = pbox[pos]
                tbox_pos = tbox[pos]
                box_loss = ciou_loss_xyxy(pbox_pos, tbox_pos).mean()
                total_box = total_box + box_loss

                # class loss
                cls_loss = self.bce(pcls_logit[pos], tcls[pos]).mean()
                total_cls = total_cls + cls_loss

        loss = self.lambda_box * total_box + self.lambda_obj * total_obj + self.lambda_cls * total_cls
        return loss, {"box": total_box.detach(), "obj": total_obj.detach(), "cls": total_cls.detach()}