set run name manually

Author: ekarais

scripts/templates/train_etnn.sh.j2

@@ -3,6 +3,7 @@
 # Training script for {{ experiment_name }}
 
 # DEFINE EXP ARGUMENTS
+EXP_NAME={{ experiment_name }}
 LIFTERS=({{ lifters | join(' ') }})
 DIM={{ dim }}
 VISIBLE_DIMS=({{ visible_dims | join(' ') }})
@@ -76,6 +77,7 @@ do
                            --splits "$SPLITS" \
                            --normalize_invariants \
                            --clip_gradient \
+                           --run_name "${EXP_NAME} ${TARGET_NAME}" \
                            --checkpoint_dir "$CHECKPOINT_DIR" &
 
     # Wait for 30 minutes before moving to the next iteration

src/main_qm9.py

@@ -118,9 +118,9 @@ def main(args):
             resume="must",
         )
     else:
-        wandb.init(project="QM9-clean-experiments")
+        run_name = args.run_name
+        wandb.init(project="QM9-clean-experiments", name=run_name)
         run_id = wandb.run.id
-        run_name = wandb.run.name
 
     for epoch in tqdm(range(start_epoch, args.epochs)):
         epoch_start_time, epoch_mae_train, epoch_mae_val = time.time(), 0, 0
@@ -290,6 +290,9 @@ def main(args):
     parser.add_argument("--num_samples", type=int, default=None, help="num samples to to train on")
     parser.add_argument("--splits", type=str, default="egnn", help="split type")
 
+    # wandb arguments
+    parser.add_argument("--run_name", type=str, default=None, help="run name")
+
     parsed_args = parser.parse_args()
     parsed_args = parser_utils.add_common_derived_arguments(parsed_args)
     parsed_args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")