1

config.py

@@ -8,7 +8,7 @@ INPUT_IMAGE_SIZE = 512
 # 5: sin(direction)
 NUM_FEATURE_MAP_CHANNEL = 6
 # image_size / 2^5 = 512 / 32 = 16
-FEATURE_MAP_SIZE = 16
+FEATURE_MAP_SIZE = 80
 # Threshold used to filter marking points too close to image boundary
 BOUNDARY_THRESH = 0.05
 
@@ -36,7 +36,7 @@ CONFID_THRESH_FOR_POINT = 0.11676871
 
 def add_common_arguments(parser):
     """Add common arguments for training and inference."""
-    parser.add_argument('--detector_weights',
+    parser.add_argument('--detector_weights', default='',
                         help="The weights of pretrained detector.")
     parser.add_argument('--depth_factor', type=int, default=32,
                         help="Depth factor.")
@@ -49,20 +49,24 @@ def add_common_arguments(parser):
 def get_parser_for_training():
     """Return argument parser for training."""
     parser = argparse.ArgumentParser()
-    parser.add_argument('--dataset_directory', required=True,
+    parser.add_argument('--dataset_directory', default='/home/thsw/ssd/zjc/AngularPoint1130/train',
                         help="The location of dataset.")
+    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml',
+                        help='model.yaml path')
     parser.add_argument('--optimizer_weights',
                         help="The weights of optimizer.")
-    parser.add_argument('--batch_size', type=int, default=16,
+    parser.add_argument('--batch_size', type=int, default=56,
                         help="Batch size.")
-    parser.add_argument('--data_loading_workers', type=int, default=0,
+    parser.add_argument('--data_loading_workers', type=int, default=8,
                         help="Number of workers for data loading.")
-    parser.add_argument('--num_epochs', type=int, default=500,
+    parser.add_argument('--num_epochs', type=int, default=300,
                         help="Number of epochs to train for.")
     parser.add_argument('--lr', type=float, default=1e-4,
                         help="The learning rate of back propagation.")
     parser.add_argument('--enable_visdom', action='store_true',
                         help="Enable Visdom to visualize training progress")
+    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml',
+                        help='hyperparameters path')
     add_common_arguments(parser)
     return parser
 

train.py

@@ -1,12 +1,20 @@
 """Train directional marking point detector."""
 import math
 import random
+
+import numpy as np
 import torch
+import yaml
+from torch import nn
 from torch.utils.data import DataLoader
 import config
 import data
 import util
 from model import DirectionalPointDetector
+from models.yolo import Model
+
+# import os
+# os.environ['CUDA_VISIBLE_DEVICES'] = '1'
 
 
 def plot_prediction(logger, image, marking_points, prediction):
@@ -51,19 +59,56 @@ def generate_objective(marking_points_batch, device):
     return objective, gradient
 
 
+# class FocalLoss(nn.Module):
+#     # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
+#     def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
+#         super(FocalLoss, self).__init__()
+#         self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
+#         self.gamma = gamma
+#         self.alpha = alpha
+#         self.reduction = loss_fcn.reduction
+#         self.loss_fcn.reduction = 'none'  # required to apply FL to each element
+#
+#     def forward(self, pred, true):
+#         loss = self.loss_fcn(pred, true)
+#         # p_t = torch.exp(-loss)
+#         # loss *= self.alpha * (1.000001 - p_t) ** self.gamma  # non-zero power for gradient stability
+#
+#         # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
+#         pred_prob = torch.sigmoid(pred)  # prob from logits
+#         p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
+#         alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
+#         modulating_factor = (1.0 - p_t) ** self.gamma
+#         loss *= alpha_factor * modulating_factor
+#
+#         if self.reduction == 'mean':
+#             return loss.mean()
+#         elif self.reduction == 'sum':
+#             return loss.sum()
+#         else:  # 'none'
+#             return loss
+
+
 def train_detector(args):
     """Train directional point detector."""
     args.cuda = not args.disable_cuda and torch.cuda.is_available()
     device = torch.device('cuda:' + str(args.gpu_id) if args.cuda else 'cpu')
     torch.set_grad_enabled(True)
 
-    dp_detector = DirectionalPointDetector(
-        3, args.depth_factor, config.NUM_FEATURE_MAP_CHANNEL).to(device)
+    # dp_detector = DirectionalPointDetector(
+    #     3, args.depth_factor, config.NUM_FEATURE_MAP_CHANNEL).to(device)
+    # if args.detector_weights:
+    #     print("Loading weights: %s" % args.detector_weights)
+    #     dp_detector.load_state_dict(torch.load(args.detector_weights))
+    # dp_detector.train()
+
+    with open(args.hyp) as f:
+        hyp = yaml.load(f, Loader=yaml.SafeLoader)
+    dp_detector = Model(args.cfg, ch=3, anchors=hyp.get('anchors')).to(device)
     if args.detector_weights:
         print("Loading weights: %s" % args.detector_weights)
         dp_detector.load_state_dict(torch.load(args.detector_weights))
     dp_detector.train()
-
     optimizer = torch.optim.Adam(dp_detector.parameters(), lr=args.lr)
     if args.optimizer_weights:
         print("Loading weights: %s" % args.optimizer_weights)
@@ -73,16 +118,27 @@ def train_detector(args):
     data_loader = DataLoader(data.ParkingSlotDataset(args.dataset_directory),
                              batch_size=args.batch_size, shuffle=True,
                              num_workers=args.data_loading_workers,
+                             pin_memory=True,
                              collate_fn=lambda x: list(zip(*x)))
 
+    # BCEobj = nn.BCEWithLogitsLoss(reduction='none', pos_weight=torch.tensor([hyp['obj_pw']], device=device))
+
+    # # Focal loss
+    # g = hyp['fl_gamma']    # focal loss gamma
+    # if g > 0:
+    #     BCEobj = FocalLoss(BCEobj, g)
+
     for epoch_idx in range(args.num_epochs):
         for iter_idx, (images, marking_points) in enumerate(data_loader):
             images = torch.stack(images).to(device)
-
+            # images = torch.from_numpy(np.stack(images, axis=0)).to(device).permute(0, 3, 1, 2)
             optimizer.zero_grad()
             prediction = dp_detector(images)
             objective, gradient = generate_objective(marking_points, device)
+            # lobj = BCEobj(prediction[:, 0, ...], objective[:, 0, ...])
             loss = (prediction - objective) ** 2
+            # lobj = torch.unsqueeze(lobj, 1)
+            # loss = torch.cat((lobj, l_sxycs), 1)
             loss.backward(gradient)
             optimizer.step()