ManjunathReddy/Yolo3_from_scratch

1"""2Implementation of Yolo Loss Function similar to the one in Yolov3 paper,3the difference from what I can tell is I use CrossEntropy for the classes4instead of BinaryCrossEntropy.5"""6import random7import torch8import torch.nn as nn9 10from src.utils_rh import intersection_over_union11 12 13class YoloLoss(nn.Module):14    def __init__(self):15        super().__init__()16        self.mse = nn.MSELoss()17        self.bce = nn.BCEWithLogitsLoss()18        self.entropy = nn.CrossEntropyLoss()19        self.sigmoid = nn.Sigmoid()20 21        # Constants signifying how much to pay for each respective part of the loss22        self.lambda_class = 123        self.lambda_noobj = 1024        self.lambda_obj = 125        self.lambda_box = 1026 27    def forward(self, predictions, target, anchors):28        # Check where obj and noobj (we ignore if target == -1)29        obj = target[..., 0] == 1  # in paper this is Iobj_i30        noobj = target[..., 0] == 0  # in paper this is Inoobj_i31 32        # ======================= #33        #   FOR NO OBJECT LOSS    #34        # ======================= #35 36        no_object_loss = self.bce(37            (predictions[..., 0:1][noobj]), (target[..., 0:1][noobj]),38        )39 40        # ==================== #41        #   FOR OBJECT LOSS    #42        # ==================== #43 44        anchors = anchors.reshape(1, 3, 1, 1, 2)45        box_preds = torch.cat([self.sigmoid(predictions[..., 1:3]), torch.exp(predictions[..., 3:5]) * anchors], dim=-1)46        ious = intersection_over_union(box_preds[obj], target[..., 1:5][obj]).detach()47        object_loss = self.mse(self.sigmoid(predictions[..., 0:1][obj]), ious * target[..., 0:1][obj])48 49        # ======================== #50        #   FOR BOX COORDINATES    #51        # ======================== #52 53        predictions[..., 1:3] = self.sigmoid(predictions[..., 1:3])  # x,y coordinates54        target[..., 3:5] = torch.log(55            (1e-16 + target[..., 3:5] / anchors)56        )  # width, height coordinates57        box_loss = self.mse(predictions[..., 1:5][obj], target[..., 1:5][obj])58 59        # ================== #60        #   FOR CLASS LOSS   #61        # ================== #62 63        class_loss = self.entropy(64            (predictions[..., 5:][obj]), (target[..., 5][obj].long()),65        )66 67        #print("__________________________________")68        #print(self.lambda_box * box_loss)69        #print(self.lambda_obj * object_loss)70        #print(self.lambda_noobj * no_object_loss)71        #print(self.lambda_class * class_loss)72        #print("\n")73 74        return (75            self.lambda_box * box_loss76            + self.lambda_obj * object_loss77            + self.lambda_noobj * no_object_loss78            + self.lambda_class * class_loss79        )