YOLO loss

객체 검출을 진행하는 AI 모델에는 2가지의 방법이 있다.

<1-stage detector, 2-stage detector>

1-stage 방식의 대표적인 모델은 YOLO 시리즈와 Retina-Net, SSD 등이 있다.

1-stage 방식이란 Regional Proposal과 Classification이 CNN을 통해 동시에 이루어지는 방식으로 Convolution Layer을 통해 Feature Maps가 생성되면 Output으로 Multi-Class Classification과 Bounding Box Regression을 출력한다.

위 방식의 구조는 Anchor Boxes(앵커 박스)를 찾게되는데 Anchor box란 중심 좌표를 기준으로 여러 크기와 비율을 가지고 생성된 영역이다.

Anchor box 관련 글

돌아와 YOLO 시리즈에서 loss값을 계산하여 학습시키는 모듈은 ultralytics/nn/tasks.py 으로 해당 모듈은 기본적으로 다른 Detection, Segmentation, Poss 등의 작업에 대한 기초 클래스 역할을 하는 BaseModel이 있다.

class BaseModel 은 모델의 기초로서 YOLOv8의 공통적인 레이어 구성과 파라미터 초기화를 초기화하여 이후에 이를 상속받는 모델들에 있어 기초 토대를 만들어준다.

기초 토대를 구성해 준다는 것은 손실 값 계산, 모델의 가중치 로드, 순방향 패스의 수행에 있어 입력값과 출력값 그 사이에 필요한 파라미터(특성 맵을 저장할지에 대한 여부)등 이 있다.

class DetectionModel(BaseModel)
        # Define model
        ch = self.yaml["ch"] = self.yaml.get("ch", ch)  # input channels
        if nc and nc != self.yaml["nc"]:
            LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
            self.yaml["nc"] = nc  # override YAML value
        self.model, self.save = parse_model(deepcopy(self.yaml), ch=ch, verbose=verbose)  # model, savelist
        self.names = {i: f"{i}" for i in range(self.yaml["nc"])}  # default names dict
        self.inplace = self.yaml.get("inplace", True)
        self.end2end = getattr(self.model[-1], "end2end", False)

def loss(self, batch, preds=None):
        """
        Compute the loss for the given batch of data.
        Args:
            batch (dict): Dictionary containing image and label data.
            preds (torch.Tensor, optional): Precomputed model predictions. Defaults to None.
        Returns:
            (tuple): A tuple containing the total loss and main three losses in a tensor.
        """
        if not hasattr(self, "criterion"):
            self.criterion = self.init_criterion()
 
        img = batch["img"]
        # NOTE: preprocess gt_bbox and gt_labels to list.
        bs = len(img)
        batch_idx = batch["batch_idx"]
        gt_groups = [(batch_idx == i).sum().item() for i in range(bs)]
        targets = {
            "cls": batch["cls"].to(img.device, dtype=torch.long).view(-1),
            "loc": batch["loc"].to(img.device, dtype=torch.long).view(-1),
            "action": batch["action"].to(img.device, dtype=torch.long).view(-1),
            "bboxes": batch["bboxes"].to(device=img.device),
            "batch_idx": batch_idx.to(img.device, dtype=torch.long).view(-1),
            "gt_groups": gt_groups,
        }
 
        preds = self.predict(img, batch=targets) if preds is None else preds
        dec_bboxes, dec_scores, enc_bboxes, enc_scores, dn_meta = preds if self.training else preds[1]
        if dn_meta is None:
            dn_bboxes, dn_scores = None, None
        else:
            dn_bboxes, dec_bboxes = torch.split(dec_bboxes, dn_meta["dn_num_split"], dim=2)
            dn_scores, dec_scores = torch.split(dec_scores, dn_meta["dn_num_split"], dim=2)
 
        dec_bboxes = torch.cat([enc_bboxes.unsqueeze(0), dec_bboxes])  # (7, bs, 300, 4)
        dec_scores = torch.cat([enc_scores.unsqueeze(0), dec_scores])
 
        loss = self.criterion(
            (dec_bboxes, dec_scores), targets, dn_bboxes=dn_bboxes, dn_scores=dn_scores, dn_meta=dn_meta
        )
        # NOTE: There are like 12 losses in RTDETR, backward with all losses but only show the main three losses.
        return sum(loss.values()), torch.as_tensor(
            [loss[k].detach() for k in ["loss_giou", "loss_class", "loss_bbox"]], device=img.device
        )

참고할 사이트

ultralytics/models/yolo/detect/train.py, val.py, predict.py

해당 모듈에서는 모델이 학습할 때 쓰이는 loss값이 어떻게 설정되고 계산하는지 포함.

손실 항목은 self.loss_names에 저장.

self.loss_names = "box_loss", "cls_loss", "dfl_loss"

https://github.com/kCW-tb/complex_detection/blob/main/ultralytics/models/yolo/detect/train.py