Working on loss

Author: odedgolden

scripts/train.py

@@ -1,6 +1,7 @@
 import torch
 from torch.utils.data import DataLoader
 
+from src.YOLOLoss import YOLOLoss
 from src.YOLOv3 import YOLOv3
 from src.YOLOv3Dataset import YOLOv3Dataset
 
@@ -11,4 +12,6 @@
 # model.load_state_dict(torch.load("weights.pt"))
 
 train_dataset = YOLOv3Dataset()
-dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)
+
+loss = YOLOLoss()

src/YOLOLoss.py

@@ -12,9 +12,9 @@ def __init__(self):
         self.mse_loss = nn.MSELoss()  # For Bounding Box Prediction
         self.bce_loss = nn.BCELoss()  # For Class Prediction
 
-    def forward(self, x, y):
-        # self.metrics, total_loss = self.calculate_metrics(pred_boxes=pred_boxes, pred_cls=pred_cls, targets=y, x=x_0,
-        #                                                   y=y_0, w=w, h=h, pred_conf=pred_conf)
+    def forward(self, predictions, ground_truth):
+        self.metrics, total_loss = self.calculate_metrics(pred_boxes=pred_boxes, pred_cls=pred_cls, targets=y, x=x_0,
+                                                          y=y_0, w=w, h=h, pred_conf=pred_conf)
 
         return
 

src/YOLOv3.py

@@ -1,6 +1,5 @@
 import torch.nn as nn
 
-from src.YOLODummy import YOLODummy
 from src.YOLOLayer import YOLOLayer
 from src.YOLOModule import YOLOModule
 

src/YOLOv3Dataset.py

@@ -1,5 +1,5 @@
 import torch
-import torch.nn as nn
+from torch.utils.data import Dataset
 from torchvision import transforms
 import torch.nn.functional as F
 
@@ -10,7 +10,7 @@
 import numpy as np
 
 
-class YOLOv3Dataset(nn.Module):
+class YOLOv3Dataset(Dataset):
 
     def __init__(self, images_list_path="../data/images_list.txt", image_size=416, should_augment=True, transform=None):
         super(YOLOv3Dataset, self).__init__()