marvin/CNN_Website
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

changed yolo loss

Browse Source
This commit is contained in:
marvin
2026-05-22 05:57:40 +02:00
parent e412738d3e
commit b6a4dd5fa2
5 changed files with 40 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/src/yolo/yolo_loss.py
+7 -4
View File
@@ -12,13 +12,16 @@ class YoloLoss(nn.Module):
target_boxes = targets[..., :4]
target_conf = targets[..., 4]
target_classes = targets[..., 5:]
obj_mask = targets[..., 4] == 1
noobj_mask = targets[..., 4] == 0
box_loss = lambda_coord * torch.mean((pred_boxes - target_boxes) ** 2)
box_loss = lambda_coord * torch.mean((pred_boxes[obj_mask] - target_boxes[obj_mask]) ** 2)
obj_loss = torch.mean((pred_conf[target_conf == 1] - target_conf[target_conf == 1]) ** 2)
noobj_loss = lambda_noobj * torch.mean((pred_conf[target_conf == 0]) ** 2)
obj_loss = torch.mean((pred_conf[obj_mask] - target_conf[obj_mask]) ** 2)
noobj_loss = lambda_noobj * torch.mean((pred_conf[noobj_mask]) ** 2)
class_loss = torch.mean((pred_classes[target_conf == 1] - target_classes[target_conf == 1]) ** 2)
class_loss = torch.mean((pred_classes[obj_mask] - target_classes[obj_mask]) ** 2)
total_loss = box_loss + obj_loss + noobj_loss + class_loss
return total_loss
api/src/yolo/yolo_model.py
+2 -1
View File
@@ -30,7 +30,8 @@ class Yolo_model(nn.Module):
nn.BatchNorm2d(c_hidden),
nn.LeakyReLU(),
nn.Flatten(),
nn.Linear(in_features=img_size*img_size*c_hidden, out_features=grid*grid*(boxes*5+labels))
nn.Linear(in_features=img_size*img_size*c_hidden, out_features=grid*grid*(boxes*5+labels)),
nn.Dropout(0.3)
)
def forward(self, x):