Add group to wandb config

talmolab · Jan 15, 2025 · 7782681 · 7782681
1 parent 0ac9fa4
commit 7782681
Show file tree

Hide file tree

Showing 11 changed files with 10 additions and 111 deletions.
diff --git a/docs/config.md b/docs/config.md
@@ -167,6 +167,7 @@ The config file has three main sections:
         - `api_key`: (str) API key. The API key is masked when saved to config files.
         - `wandb_mode`: (str) "offline" if only local logging is required. Default: "None".
         - `prv_runid`: (str) Previous run ID if training should be resumed from a previous ckpt. *Default*: `None`.
+        - `group`: (str) Group name for the run.
     - `optimizer_name`: (str) Optimizer to be used. One of ["Adam", "AdamW"].
     - `optimizer`
         - `lr`: (float) Learning rate of type float. *Default*: 1e-3

diff --git a/docs/config_centroid.yaml b/docs/config_centroid.yaml
@@ -120,6 +120,7 @@ trainer_config:
     wandb_mode: ''
     api_key: ''
     prv_runid:
+    group:
   optimizer_name: Adam
   optimizer:
     lr: 0.0001

diff --git a/initial_config.yaml b/initial_config.yaml
diff --git a/sleap_nn/training/model_trainer.py b/sleap_nn/training/model_trainer.py
@@ -639,6 +639,7 @@ def train(
                 name=wandb_config.name,
                 save_dir=self.dir_path,
                 id=self.config.trainer_config.wandb.prv_runid,
+                group=self.config.trainer_config.wandb.group,
             )
             logger.append(wandb_logger)
 

diff --git a/tests/assets/minimal_instance/initial_config.yaml b/tests/assets/minimal_instance/initial_config.yaml
@@ -75,6 +75,7 @@ trainer_config:
     wandb_mode: ''
     api_key: ''
     prv_runid:
+    group:
   optimizer_name: Adam
   optimizer:
     lr: 0.0001

diff --git a/tests/assets/minimal_instance/training_config.yaml b/tests/assets/minimal_instance/training_config.yaml
@@ -88,6 +88,7 @@ trainer_config:
     wandb_mode: ''
     api_key: ''
     prv_runid: null
+    group:
   optimizer_name: Adam
   optimizer:
     lr: 0.0001

diff --git a/tests/assets/minimal_instance_bottomup/initial_config.yaml b/tests/assets/minimal_instance_bottomup/initial_config.yaml
@@ -76,6 +76,7 @@ trainer_config:
     wandb_mode: ''
     api_key: ''
     prv_runid:
+    group:
   optimizer_name: Adam
   optimizer:
     lr: 0.0001

diff --git a/tests/assets/minimal_instance_bottomup/training_config.yaml b/tests/assets/minimal_instance_bottomup/training_config.yaml
@@ -91,6 +91,7 @@ trainer_config:
     wandb_mode: ''
     api_key: ''
     prv_runid:
+    group:
   optimizer_name: Adam
   optimizer:
     lr: 0.0001

diff --git a/tests/assets/minimal_instance_centroid/initial_config.yaml b/tests/assets/minimal_instance_centroid/initial_config.yaml
@@ -70,6 +70,7 @@ trainer_config:
     wandb_mode: ''
     api_key: ''
     prv_runid:
+    group:
   optimizer_name: Adam
   optimizer:
     lr: 0.0001

diff --git a/tests/assets/minimal_instance_centroid/training_config.yaml b/tests/assets/minimal_instance_centroid/training_config.yaml
@@ -81,6 +81,7 @@ trainer_config:
     wandb_mode: ''
     api_key: ''
     prv_runid:
+    group:
   optimizer_name: Adam
   optimizer:
     lr: 0.0001

diff --git a/tests/fixtures/datasets.py b/tests/fixtures/datasets.py
@@ -145,13 +145,7 @@ def config(sleap_data_dir):
                     "wandb_mode": "offline",
                     "api_key": "",
                     "prv_runid": None,
-                    "log_params": [
-                        "trainer_config.optimizer_name",
-                        "trainer_config.optimizer.amsgrad",
-                        "trainer_config.optimizer.lr",
-                        "model_config.backbone_type",
-                        "model_config.init_weights",
-                    ],
+                    "group": None,
                 },
                 "optimizer_name": "Adam",
                 "optimizer": {"lr": 0.0001, "amsgrad": False},