#23: 5 class loss modified

 1. torch tensor type change to TorchFloat.cuda() in one_hot function from utils.py
 2. stacking objectness class map for matching tensor dimension in detection_loss_4_small_yolo from yolov1_small.py

utilities/utils.py

@@ -2,6 +2,7 @@
 import torch
 import imgaug as ia
 import torchvision.transforms as transforms
+import torch
 
 import matplotlib.pyplot as plt
 from PIL import Image, ImageDraw
@@ -20,7 +21,10 @@ def one_hot(output , label):
 
                 dst[k][i][j][int(label[k][i][j])] = 1.
 
-    return torch.from_numpy(dst)
+    result = torch.from_numpy(dst)
+    result = result.type(torch.FloatTensor).cuda()
+
+    return result
 
 # visdom function
 
@@ -51,19 +55,6 @@ def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
     if is_best:
         shutil.copyfile(filename, 'model_best.pth.tar')
 
-def one_hot(output , label):
-
-    label = label.cpu().data.numpy()
-    b, s1, s2, c = output.shape
-    dst = np.zeros([b,s1,s2,c], dtype=np.float32)
-
-    for k in range(b):
-        for i in range(s1):
-            for j in range(s2):
-
-                dst[k][i][j][int(label[k][i][j])] = 1.
-
-    return torch.from_numpy(dst)
 
 def CvtCoordsXXYY2XYWH(image_width, image_height, xmin, xmax, ymin, ymax):
     #calculate bbox_center

yolov1_small.py

@@ -187,33 +187,40 @@ def forward(self, x):
         return out
 
 def detection_loss_4_small_yolo(output, target):
+    from utilities.utils import one_hot
+
     # hyper parameter
 
     lambda_coord = 5
     lambda_noobj = 0.5
 
     # check batch size
     b, _, _, _ = target.shape
+    _, _, _, n = output.shape
+
+    # calc number of class
+    num_of_cls = n-5
 
     # class loss
     MSE_criterion = nn.MSELoss()
 
     # output tensor slice
+    # output tensor shape is [batch, 7, 7, 5 + classes]
     objness1_output = output[:, :, :, 0]
     x_offset1_output = output[:, :, :, 1]
     y_offset1_output = output[:, :, :, 2]
     width_ratio1_output = output[:, :, :, 3]
     height_ratio1_output = output[:, :, :, 4]
-    class_output = output[:, :, :, 5]
+    class_output = output[:, :, :, 5:]
 
     # label tensor slice
     objness_label = target[:, :, :, 0]
-    class_label = target[:, :, :, 1]
+    class_label = one_hot(class_output, target[:, :, :, 1])
     x_offset_label = target[:, :, :, 2]
     y_offset_label = target[:, :, :, 3]
     width_ratio_label = target[:, :, :, 4]
     height_ratio_label = target[:, :, :, 5]
-    
+
     noobjness_label = torch.neg(torch.add(objness_label, -1))
 
     obj_coord1_loss = lambda_coord * \
@@ -226,12 +233,14 @@ def detection_loss_4_small_yolo(output, target):
                                (torch.pow(width_ratio1_output - torch.sqrt(width_ratio_label), 2) +
                                 torch.pow(height_ratio1_output - torch.sqrt(height_ratio_label), 2)))
 
-    obj_class_loss = torch.sum(objness_label * torch.pow(class_output - class_label, 2))
+    objectness_cls_map = torch.stack((objness_label,objness_label,objness_label,objness_label,objness_label), 3)
+    obj_class_loss = torch.sum(objectness_cls_map * torch.pow(class_output - class_label, 2))
 
     noobjness1_loss = lambda_noobj * torch.sum(noobjness_label * torch.pow(objness1_output - objness_label, 2))
 
     objness1_loss = torch.sum(objness_label * torch.pow(objness1_output - objness_label, 2))
 
-    total_loss = (obj_coord1_loss + obj_size1_loss + noobjness1_loss  + objness1_loss + obj_class_loss) / b
+    total_loss = (obj_coord1_loss + obj_size1_loss + noobjness1_loss  + objness1_loss + obj_class_loss)
+    total_loss = total_loss / b
 
     return total_loss, obj_coord1_loss / b, obj_size1_loss/ b, obj_class_loss/ b, noobjness1_loss/ b, objness1_loss/ b