Merge branch 'master' into data_parallel

1 files changed, 12 insertions, 8 deletions

diff --git a/models/model.py b/models/model.py
index a8e0316..1f8ae23 100644
--- a/models/model.py
+++ b/models/model.py
@@ -187,11 +187,14 @@ class Model:
         ], **optim_hp)
         sched_final_gamma = sched_hp.get('final_gamma', 0.001)
         sched_start_step = sched_hp.get('start_step', 15_000)
+        all_step = self.total_iter - sched_start_step
 
         def lr_lambda(epoch):
-            passed_step = epoch - sched_start_step
-            all_step = self.total_iter - sched_start_step
-            return sched_final_gamma ** (passed_step / all_step)
+            if epoch > sched_start_step:
+                passed_step = epoch - sched_start_step
+                return sched_final_gamma ** (passed_step / all_step)
+            else:
+                return 1
         self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=[
             lr_lambda, lr_lambda, lr_lambda, lr_lambda
         ])
@@ -286,13 +289,14 @@ class Model:
                 'Embedding/PartNet norm', mean_pa_norm,
                 self.k, self.pr * self.k, self.curr_iter
             )
+            # Learning rate
+            lrs = self.scheduler.get_last_lr()
+            # Write learning rates
+            self.writer.add_scalar(
+                'Learning rate', lrs[0], self.curr_iter
+            )
 
             if self.curr_iter % 100 == 0:
-                lrs = self.scheduler.get_last_lr()
-                # Write learning rates
-                self.writer.add_scalar(
-                    'Learning rate', lrs[0], self.curr_iter
-                )
                 # Write disentangled images
                 if self.image_log_on:
                     i_a, i_c, i_p = images