From af853f604c69d8d138efe9a377c0d2e7a0d80cd8 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 21 Feb 2019 15:57:18 +0100
Subject: [PATCH] updates

---
 models.py |  2 +-
 train.py  | 42 +++++++++++++++++++++---------------------
 2 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/models.py b/models.py
index 2386b222..7a659404 100755
--- a/models.py
+++ b/models.py
@@ -295,7 +295,7 @@ def load_darknet_weights(self, weights, cutoff=-1):
     if weights_file == 'darknet53.conv.74':
         cutoff = 75
     elif weights_file == 'yolov3-tiny.conv.15':
-        cutoff = 16
+        cutoff = 15
 
     # Open the weights file
     fp = open(weights, 'rb')
diff --git a/train.py b/train.py
index 1239fcb5..2fa057d6 100644
--- a/train.py
+++ b/train.py
@@ -15,11 +15,13 @@ def train(
         epochs=100,
         batch_size=16,
         accumulated_batches=1,
-        weights='weights',
         multi_scale=False,
         freeze_backbone=True,
         var=0,
 ):
+    weights = 'weights' + os.sep
+    latest = weights + 'latest.pt'
+    best = weights + 'best.pt'
     device = torch_utils.select_device()
 
     if multi_scale:  # pass maximum multi_scale size
@@ -27,9 +29,6 @@ def train(
     else:
         torch.backends.cudnn.benchmark = True  # unsuitable for multiscale
 
-    latest = os.path.join(weights, 'latest.pt')
-    best = os.path.join(weights, 'best.pt')
-
     # Configure run
     train_path = parse_data_cfg(data_cfg)['train']
 
@@ -40,6 +39,7 @@ def train(
     dataloader = LoadImagesAndLabels(train_path, batch_size, img_size, multi_scale=multi_scale, augment=True)
 
     lr0 = 0.001
+    cutoff = -1  # backbone reaches to cutoff layer
     if resume:
         checkpoint = torch.load(latest, map_location='cpu')
 
@@ -69,8 +69,13 @@ def train(
         start_epoch = 0
         best_loss = float('inf')
 
-        # Initialize model with darknet53 weights (optional)
-        load_darknet_weights(model, os.path.join(weights, 'darknet53.conv.74'))
+        # Initialize model with backbone (optional)
+        if cfg.endswith('yolov3.cfg'):
+            load_darknet_weights(model, weights + 'darknet53.conv.74')
+            cutoff = 75
+        elif cfg.endswith('yolov3-tiny.cfg'):
+            load_darknet_weights(model, weights + 'yolov3-tiny.conv.15')
+            cutoff = 15
 
         # if torch.cuda.device_count() > 1:
         #     model = nn.DataParallel(model)
@@ -102,15 +107,10 @@ def train(
             g['lr'] = lr
 
         # Freeze darknet53.conv.74 for first epoch
-        if freeze_backbone:
-            if epoch == 0:
-                for i, (name, p) in enumerate(model.named_parameters()):
-                    if int(name.split('.')[1]) < 75:  # if layer < 75
-                        p.requires_grad = False
-            elif epoch == 1:
-                for i, (name, p) in enumerate(model.named_parameters()):
-                    if int(name.split('.')[1]) < 75:  # if layer < 75
-                        p.requires_grad = True
+        if freeze_backbone and (epoch < 2):
+            for i, (name, p) in enumerate(model.named_parameters()):
+                if int(name.split('.')[1]) < cutoff:  # if layer < 75
+                    p.requires_grad = False if (epoch == 0) else True
 
         ui = -1
         rloss = defaultdict(float)  # running loss
@@ -140,9 +140,11 @@ def train(
                 rloss[key] = (rloss[key] * ui + val) / (ui + 1)
 
             s = ('%8s%12s' + '%10.3g' * 7) % (
-                '%g/%g' % (epoch, epochs - 1), '%g/%g' % (i, len(dataloader) - 1), rloss['xy'],
-                rloss['wh'], rloss['conf'], rloss['cls'],
-                rloss['loss'], model.losses['nT'], time.time() - t0)
+                '%g/%g' % (epoch, epochs - 1),
+                '%g/%g' % (i, len(dataloader) - 1),
+                rloss['xy'], rloss['wh'], rloss['conf'],
+                rloss['cls'], rloss['loss'],
+                model.losses['nT'], time.time() - t0)
             t0 = time.time()
             print(s)
 
@@ -164,7 +166,7 @@ def train(
 
         # Save backup weights every 5 epochs (optional)
         # if (epoch > 0) & (epoch % 5 == 0):
-        #     os.system('cp ' + latest + ' ' + os.path.join(weights, 'backup{}.pt'.format(epoch)))
+        #     os.system('cp ' + latest + ' ' + weights + 'backup{}.pt'.format(epoch)))
 
         # Calculate mAP
         with torch.no_grad():
@@ -184,7 +186,6 @@ if __name__ == '__main__':
     parser.add_argument('--data-cfg', type=str, default='cfg/coco.data', help='coco.data file path')
     parser.add_argument('--multi-scale', action='store_true', help='random image sizes per batch 320 - 608')
     parser.add_argument('--img-size', type=int, default=32 * 13, help='pixels')
-    parser.add_argument('--weights', type=str, default='weights', help='path to store weights')
     parser.add_argument('--resume', action='store_true', help='resume training flag')
     parser.add_argument('--var', type=float, default=0, help='test variable')
     opt = parser.parse_args()
@@ -200,7 +201,6 @@ if __name__ == '__main__':
         epochs=opt.epochs,
         batch_size=opt.batch_size,
         accumulated_batches=opt.accumulated_batches,
-        weights=opt.weights,
         multi_scale=opt.multi_scale,
         var=opt.var,
     )