move scheduler.step() to train()

digitalbrain79 · Mar 27, 2019 · 1e9d466 · 1e9d466
1 parent 4cd20b5
commit 1e9d466
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/torchreid/engine/engine.py b/torchreid/engine/engine.py
@@ -99,9 +99,6 @@ def run(self, save_dir='log', max_epoch=0, start_epoch=0, fixbase_epoch=0, open_
         for epoch in range(start_epoch, max_epoch):
             self.train(epoch, trainloader, print_freq=print_freq)
 
-            if self.scheduler is not None:
-                self.scheduler.step()
-
             if (epoch+1)>start_eval and eval_freq>0 and (epoch+1)%eval_freq==0 and (epoch+1)!=max_epoch:
                 rank1 = self.test(
                     epoch,

diff --git a/torchreid/engine/image/softmax.py b/torchreid/engine/image/softmax.py
@@ -126,4 +126,7 @@ def train(self, epoch, trainloader, fixbase=False, open_layers=None, print_freq=
                       loss=losses,
                       acc=accs))
 
-            end = time.time()
+            end = time.time()
+
+        if self.scheduler is not None:
+            self.scheduler.step()
diff --git a/torchreid/engine/image/triplet.py b/torchreid/engine/image/triplet.py
@@ -143,4 +143,7 @@ def train(self, epoch, trainloader, fixbase=False, open_layers=None, print_freq=
                       loss_x=losses_x,
                       acc=accs))
 
-            end = time.time()
+            end = time.time()
+
+        if self.scheduler is not None:
+            self.scheduler.step()