add lr to tensorboard

yizhou-wang · Feb 7, 2022 · b33860d · b33860d
1 parent 39ad741
commit b33860d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tools/train.py b/tools/train.py
@@ -251,6 +251,7 @@ def parse_args():
                 writer.add_scalar('loss/loss_ave', loss_ave, iter_count)
                 writer.add_scalar('time/time_load', load_time, iter_count)
                 writer.add_scalar('time/time_back', back_time, iter_count)
+                writer.add_scalar('param/param_lr', scheduler.get_last_lr()[0], iter_count)
 
                 if stacked_num is not None:
                     confmap_pred = confmap_preds[stacked_num - 1].cpu().detach().numpy()