From 1b52d55f02614353680a925508f055f8444b8e27 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marvin=20Krau=C3=9Fer?= Date: Fri, 5 Jun 2026 20:38:26 +0200 Subject: [PATCH] loss: 3.7 --- api/src/yolo/train_yolo_faces.py | 16 ++++++++-------- api/src/yolo/yolo_dataset.py | 2 +- api/src/yolo/yolo_loss.py | 6 +----- api/src/yolo/yolo_model.py | 5 +---- 4 files changed, 11 insertions(+), 18 deletions(-) diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py index 42d091d..4407992 100644 --- a/api/src/yolo/train_yolo_faces.py +++ b/api/src/yolo/train_yolo_faces.py @@ -15,7 +15,7 @@ import torch.nn.functional as F from production.yolo_model_production import convert_prediction def view_data(dataset): - dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=False) + dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=True) for images, labels in iter(dataloader): for batch in range(images.shape[0]): image = images[batch] @@ -112,16 +112,16 @@ def train_yolo(): BATCH_SIZE = 64 dataset = YoloDataset( - image_dir="data/faces_2/train", - annotation_path="data/faces_2/train/_annotations.coco.json", + image_dir="data/faces_scenery/train", + annotation_path="data/faces_scenery/train/_annotations.coco.json", img_size=IMAGE_SIZE, transform=True, grid=GRID ) dataset_valid = YoloDataset( - image_dir="data/faces_2/test", - annotation_path="data/faces_2/test/_annotations.coco.json", + image_dir="data/faces_scenery/test", + annotation_path="data/faces_scenery/test/_annotations.coco.json", img_size=IMAGE_SIZE, transform=True, grid=GRID @@ -140,7 +140,7 @@ def train_yolo(): loss_module = YoloLoss() - #view_data(dataset) + #view_data(dataset_valid) #exit() @@ -152,8 +152,8 @@ def train_yolo(): #exit() - use_webcam(GRID, IMAGE_SIZE) - exit() + #use_webcam(GRID, IMAGE_SIZE) + #exit() train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader, diff --git a/api/src/yolo/yolo_dataset.py b/api/src/yolo/yolo_dataset.py index f2a0943..732ac3f 100644 --- a/api/src/yolo/yolo_dataset.py +++ b/api/src/yolo/yolo_dataset.py @@ -48,7 +48,7 @@ class YoloDataset(Dataset): self.image_ids = list(self.coco.imgs.keys()) self.img_size = img_size self.grid = grid - self.num_classes = len(self.coco.cats)-1 + self.num_classes = len(self.coco.cats) self.toTensor = transforms.ToTensor() if transform: diff --git a/api/src/yolo/yolo_loss.py b/api/src/yolo/yolo_loss.py index 015d9dc..b88f1ac 100644 --- a/api/src/yolo/yolo_loss.py +++ b/api/src/yolo/yolo_loss.py @@ -6,7 +6,7 @@ class YoloLoss(nn.Module): def __init__(self): super(YoloLoss, self).__init__() - def forward(self, predictions, targets, lambda_coord=1, lambda_noobj=1): + def forward(self, predictions, targets): pred_boxes = predictions[..., :4] pred_conf = predictions[..., 4] pred_classes = predictions[..., 5:] @@ -27,8 +27,6 @@ class YoloLoss(nn.Module): else: box_loss = torch.tensor(0.0, device=predictions.device) - box_loss = lambda_coord * box_loss - obj_loss = F.mse_loss( pred_conf[obj_mask], @@ -42,8 +40,6 @@ class YoloLoss(nn.Module): reduction="mean" ) if noobj_mask.any() else torch.tensor(0.0, device=predictions.device) - noobj_loss = lambda_noobj * noobj_loss - class_loss = F.binary_cross_entropy_with_logits( pred_classes[obj_mask], diff --git a/api/src/yolo/yolo_model.py b/api/src/yolo/yolo_model.py index c86ae0f..7ce4329 100644 --- a/api/src/yolo/yolo_model.py +++ b/api/src/yolo/yolo_model.py @@ -80,10 +80,7 @@ class Yolo_model(nn.Module): def forward(self, x): x = self.model(x).permute(0, 2, 3, 1) - center = F.sigmoid(x[..., :2]) - size = torch.exp(x[..., 2:4]) - conf_class = F.sigmoid(x[..., 4:]) - return torch.cat([center, size, conf_class], dim=3) + return F.sigmoid(x) def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True):