1

config.py

@@ -53,11 +53,11 @@ def get_parser_for_training():
                         help="The location of dataset.")
     parser.add_argument('--optimizer_weights',
                         help="The weights of optimizer.")
-    parser.add_argument('--batch_size', type=int, default=24,
+    parser.add_argument('--batch_size', type=int, default=16,
                         help="Batch size.")
-    parser.add_argument('--data_loading_workers', type=int, default=32,
+    parser.add_argument('--data_loading_workers', type=int, default=0,
                         help="Number of workers for data loading.")
-    parser.add_argument('--num_epochs', type=int, default=20,
+    parser.add_argument('--num_epochs', type=int, default=500,
                         help="Number of epochs to train for.")
     parser.add_argument('--lr', type=float, default=1e-4,
                         help="The learning rate of back propagation.")

train.py

@@ -40,8 +40,8 @@ def generate_objective(marking_points_batch, device):
             # Makring Point Shape Regression
             objective[batch_idx, 1, row, col] = marking_point.shape
             # Offset Regression
-            objective[batch_idx, 2, row, col] = marking_point.x*16 - col
-            objective[batch_idx, 3, row, col] = marking_point.y*16 - row
+            objective[batch_idx, 2, row, col] = marking_point.x*config.FEATURE_MAP_SIZE - col
+            objective[batch_idx, 3, row, col] = marking_point.y*config.FEATURE_MAP_SIZE - row
             # Direction Regression
             direction = marking_point.direction
             objective[batch_idx, 4, row, col] = math.cos(direction)