improve heart_train.py

ChenS676 · Jun 16, 2024 · cbaaaaa · cbaaaaa
1 parent d874659
commit cbaaaaa
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/core/gcns/final_gnn_tune.py b/core/gcns/final_gnn_tune.py
@@ -287,7 +287,7 @@ def project_main(): # sourcery skip: avoid-builtin-shadow, low-code-quality
             print_logger.info(f'Num parameters: {cfg.model.params}')
 
             optimizer = create_optimizer(model, cfg)
-            scheduler = LinearDecayLR(optimizer, start_lr=0.01, end_lr=0.001, num_epochs=cfg.train.epochs)
+            scheduler = LinearDecayLR(optimizer, start_lr=cfg.optimizer.base_lr, end_lr=cfg.optimizer.base_lr/10, num_epochs=cfg.train.epochs)
 
             if cfg.train.finetune: 
                 model = init_model_from_pretrained(model, cfg.train.finetune,

diff --git a/core/graphgps/train/heart_train.py b/core/graphgps/train/heart_train.py
@@ -402,7 +402,7 @@ def train(self):
                 wandb.log({"lr": self.scheduler.get_lr()}, step=self.step)
 
             self.tensorboard_writer.add_scalar("Loss/train", loss, epoch)
-
+            self.tensorboard_writer.add_scalar("LR/train", self.scheduler.get_lr(), epoch)
             if epoch % int(self.report_step) == 0:
 
                 self.results_rank = self.merge_result_rank()