update training code

Author: hunglc007

core/yolov4.py

@@ -8,12 +8,11 @@
 import core.backbone as backbone
 from core.config import cfg
 
-
-NUM_CLASS       = len(utils.read_class_names(cfg.YOLO.CLASSES))
-STRIDES         = np.array(cfg.YOLO.STRIDES)
-IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
-XYSCALE = cfg.YOLO.XYSCALE
-ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
+# NUM_CLASS       = len(utils.read_class_names(cfg.YOLO.CLASSES))
+# STRIDES         = np.array(cfg.YOLO.STRIDES)
+# IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
+# XYSCALE = cfg.YOLO.XYSCALE
+# ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
 
 def YOLOv3(input_layer, NUM_CLASS):
     route_1, route_2, conv = backbone.darknet53(input_layer)
@@ -251,8 +250,7 @@ def bbox_giou(boxes1, boxes2):
 
     return giou
 
-# def compute_loss(pred, conv, label, bboxes, STRIDES, NUM_CLASS, IOU_LOSS_THRESH, i=0):
-def compute_loss(pred, conv, label, bboxes, i=0):
+def compute_loss(pred, conv, label, bboxes, STRIDES, NUM_CLASS, IOU_LOSS_THRESH, i=0):
     conv_shape  = tf.shape(conv)
     batch_size  = conv_shape[0]
     output_size = conv_shape[1]

train.py

@@ -43,6 +43,7 @@ def main(_argv):
         bbox_tensors = []
         for i, fm in enumerate(feature_maps):
             bbox_tensor = decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
+            bbox_tensors.append(fm)
             bbox_tensors.append(bbox_tensor)
         model = tf.keras.Model(input_layer, bbox_tensors)
     else:
@@ -51,13 +52,15 @@ def main(_argv):
             bbox_tensors = []
             for i, fm in enumerate(feature_maps):
                 bbox_tensor = decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
+                bbox_tensors.append(fm)
                 bbox_tensors.append(bbox_tensor)
             model = tf.keras.Model(input_layer, bbox_tensors)
         elif FLAGS.model == 'yolov4':
             feature_maps = YOLOv4(input_layer, NUM_CLASS)
             bbox_tensors = []
             for i, fm in enumerate(feature_maps):
                 bbox_tensor = decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
+                bbox_tensors.append(fm)
                 bbox_tensors.append(bbox_tensor)
             model = tf.keras.Model(input_layer, bbox_tensors)
 
@@ -89,8 +92,7 @@ def train_step(image_data, target):
             # optimizing process
             for i in range(3):
                 conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
-                # loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
-                loss_items = compute_loss(pred, conv, *target[i], i)
+                loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
                 giou_loss += loss_items[0]
                 conf_loss += loss_items[1]
                 prob_loss += loss_items[2]
@@ -129,7 +131,7 @@ def test_step(image_data, target):
             # optimizing process
             for i in range(3):
                 conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
-                loss_items = compute_loss(pred, conv, *target[i], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
+                loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
                 giou_loss += loss_items[0]
                 conf_loss += loss_items[1]
                 prob_loss += loss_items[2]