add lr to tensorboard

b33860d8 · Yizhou Wang · 39ad741a · b33860d8
Commit b33860d8 authored Feb 06, 2022 by Yizhou Wang
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

tools/train.py tools/train.py +1 -0

No files found.
--- a/tools/train.py
+++ b/tools/train.py
@@ -251,6 +251,7 @@ if __name__ == "__main__":
                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()