From a8e13900280fe1381f38e0e2cb009b46f8b8e498 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 17 Jan 2020 10:55:30 -0800
Subject: [PATCH] updates

---
 detect.py |  2 +-
 models.py |  2 --
 test.py   |  2 +-
 train.py  | 17 ++++++++++-------
 4 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/detect.py b/detect.py
index 3b2d17ea..2d826a49 100644
--- a/detect.py
+++ b/detect.py
@@ -25,7 +25,7 @@ def detect(save_img=False):
     if weights.endswith('.pt'):  # pytorch format
         model.load_state_dict(torch.load(weights, map_location=device)['model'])
     else:  # darknet format
-        _ = load_darknet_weights(model, weights)
+        load_darknet_weights(model, weights)
 
     # Second-stage classifier
     classify = False
diff --git a/models.py b/models.py
index ae861cec..e93171f5 100755
--- a/models.py
+++ b/models.py
@@ -351,8 +351,6 @@ def load_darknet_weights(self, weights, cutoff=-1):
             conv_layer.weight.data.copy_(conv_w)
             ptr += num_w
 
-    return cutoff
-
 
 def save_weights(self, path='model.weights', cutoff=-1):
     # Converts a PyTorch model to Darket format (*.pt to *.weights)
diff --git a/test.py b/test.py
index b0954720..bd859580 100644
--- a/test.py
+++ b/test.py
@@ -35,7 +35,7 @@ def test(cfg,
         if weights.endswith('.pt'):  # pytorch format
             model.load_state_dict(torch.load(weights, map_location=device)['model'])
         else:  # darknet format
-            _ = load_darknet_weights(model, weights)
+            load_darknet_weights(model, weights)
 
         if torch.cuda.device_count() > 1:
             model = nn.DataParallel(model)
diff --git a/train.py b/train.py
index 1dc05bf1..ad723a9c 100644
--- a/train.py
+++ b/train.py
@@ -84,12 +84,15 @@ def train():
     model = Darknet(cfg, arc=opt.arc).to(device)
 
     # Optimizer
-    pg0, pg1 = [], []  # optimizer parameter groups
+    pg0, pg1, pg2 = [], [], []  # optimizer parameter groups
     for k, v in dict(model.named_parameters()).items():
-        if 'Conv2d.weight' in k:
-            pg1 += [v]  # parameter group 1 (apply weight_decay)
+        print(k)
+        if '.bias' in k:
+            pg2 += [v]  # biases
+        elif 'Conv2d.weight' in k:
+            pg1 += [v]  # apply weight_decay
         else:
-            pg0 += [v]  # parameter group 0
+            pg0 += [v]  # all else
 
     if opt.adam:
         optimizer = optim.Adam(pg0, lr=hyp['lr0'])
@@ -97,12 +100,12 @@ def train():
     else:
         optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
     optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']})  # add pg1 with weight_decay
-    del pg0, pg1
+    optimizer.add_param_group({'params': pg2})  # add pg2
+    del pg0, pg1, pg2
 
     # https://github.com/alphadl/lookahead.pytorch
     # optimizer = torch_utils.Lookahead(optimizer, k=5, alpha=0.5)
 
-    cutoff = -1  # backbone reaches to cutoff layer
     start_epoch = 0
     best_fitness = float('inf')
     attempt_download(weights)
@@ -134,7 +137,7 @@ def train():
 
     elif len(weights) > 0:  # darknet format
         # possible weights are '*.weights', 'yolov3-tiny.conv.15',  'darknet53.conv.74' etc.
-        cutoff = load_darknet_weights(model, weights)
+        load_darknet_weights(model, weights)
 
     # Scheduler https://github.com/ultralytics/yolov3/issues/238
     # lf = lambda x: 1 - x / epochs  # linear ramp to zero