Print configuration and hyperaparameters

Display configuration and hyperparameters

Remove double import and change colors

Fixing custom models parameters print

Author: sofiene-boutaj

docs/custom_config.md

@@ -79,5 +79,5 @@ Here is a non exhaustive list of the parameters that you may want to override pe
 ### Transformation invariance
 | Name | Type | Description |
 |------|---|---|
-| task.pre_comp_emb_batch_size | int | Batch size for precomputing the embeddings. |
+| task.transformation_invariance_batch_size | int | Batch size for the transformations. |
 | task.nb_images | int | Number of images to use. |

src/thunder/benchmark.py

@@ -8,7 +8,7 @@
 
 from .datasets.utils import is_dataset_available
 from .models.utils import is_model_available, load_custom_model_from_file
-
+from .utils.utils import print_task_hyperparams
 
 def benchmark(
     model: str | Callable,
@@ -46,17 +46,19 @@ def benchmark(
     from omegaconf import OmegaConf
 
     from .utils.config import get_config
-
+    
     wandb_mode = "online" if online_wandb else "offline"
     adaptation_type = "lora" if lora else "frozen"
     ckpt_saving = "save_ckpts_all_epochs" if ckpt_save_all else "save_best_ckpt_only"
     embedding_recomputing = "recomp_embs" if recomp_embs else "no_recomp_embs"
     model_retraining = "retrain_model" if retrain_model else "no_retrain_model"
     model_name = model if isinstance(model, str) else None
+    custom_name = None
 
     if model_name and model_name.startswith("custom:"):
         model = load_custom_model_from_file(model_name.split(":")[1])
         model_name = None
+        custom_name = model.name
 
     # Get Config
     cfg = get_config(
@@ -72,6 +74,8 @@ def benchmark(
         **kwargs,
     )
 
+    print_task_hyperparams(cfg, custom_name=custom_name)
+
     if not is_dataset_available(dataset):
         from . import download_datasets
 
@@ -159,7 +163,6 @@ def run_benchmark(cfg: DictConfig, model_cls: Callable = None) -> None:
         mode=cfg.wandb.mode,
     )
     wandb_base_folder = f"|{task_type}| |{adaptation_type}| |{dataset_name}|"
-
     # Folder to save results
     res_folder = os.path.join(
         os.environ["THUNDER_BASE_DATA_FOLDER"],

src/thunder/config/task/adversarial_attack.yaml

@@ -1,7 +1,6 @@
 type: adversarial_attack
 compatible_adaptation_types: ["frozen"]
 base_embeddings_folder: ${oc.env:THUNDER_BASE_DATA_FOLDER}/embeddings/
-pre_comp_emb_batch_size: 128
 attack_batch_size: 8
 nb_attack_images: 10000
 # ----------------------------------------------------------------------

src/thunder/config/task/transformation_invariance.yaml

@@ -1,5 +1,5 @@
 type: transformation_invariance
 compatible_adaptation_types: ["frozen"]
 base_embeddings_folder: ${oc.env:THUNDER_BASE_DATA_FOLDER}/embeddings/
-pre_comp_emb_batch_size: 128
+transformation_invariance_batch_size: 64
 nb_images: 1000

src/thunder/tasks/transformation_invariance.py

@@ -143,7 +143,7 @@ def transformation_invariance(
 
     dataloader = DataLoader(
         subset_dataset,
-        batch_size=cfg.adaptation.batch_size,
+        batch_size=cfg.task.transformation_invariance_batch_size,
         shuffle=False,
         num_workers=cfg.adaptation.num_workers,
         generator=generator,

src/thunder/utils/utils.py

@@ -1,6 +1,5 @@
 import json
 import numpy as np
-from omegaconf import DictConfig
 import os
 import random
 import torch
@@ -9,6 +8,79 @@
 import h5py
 from typing import Dict, Any, Optional
 
+from typing import Optional
+from omegaconf import DictConfig, OmegaConf
+import logging
+
+def print_task_hyperparams(
+    cfg: DictConfig,
+    custom_name: Optional[str] = None
+) -> None:
+    """
+    Print dataset, model, and only the task-specific hyper-parameters
+    from a full Hydra cfg, using classic ANSI colors.
+
+    If `custom_name` is provided, it will be used instead of
+    cfg.pretrained_model.model_name.
+    """
+    RESET     = "\033[0m"
+    BOLD      = "\033[1m"
+    UNDERLINE = "\033[4m"
+    BLUE      = "\033[34m"
+    WHITE     = "\033[37m"
+    GREEN     = "\033[32m"
+    RED       = "\033[31m"
+
+    task         = cfg.task.type
+    dataset_name = cfg.dataset.dataset_name
+
+    # Safe-fetch the model name, falling back to custom_name if given
+    if custom_name is not None:
+        model_label = custom_name
+    else:
+        model_label = OmegaConf.select(
+            cfg, "pretrained_model.model_name", default="<unknown model>"
+        )
+
+    # Choose where hyperparams live
+    task_cfg = cfg.task if task not in ["linear_probing", "segmentation"] else cfg.adaptation
+
+    sep = "-" * 60
+
+    logging.info(f"\n{BOLD}{BLUE}\U0001F680  Experiment Info{RESET}")
+    print(sep)
+    print(f"{BLUE}Task   :{RESET} {WHITE}{task}{RESET}")
+    print(f"{BLUE}Dataset:{RESET} {WHITE}{dataset_name}{RESET}")
+    print(f"{BLUE}Model  :{RESET} {WHITE}{model_label}{RESET}")
+    print(sep)
+
+    print(f"\n{BOLD}{BLUE}Hyper-parameters{RESET}\n{sep}")
+
+    # Fields to skip
+    skip = {"compatible_adaptation_types", "base_embeddings_folder"}
+    if task not in ["linear_probing", "segmentation"]:
+        skip.add("type")
+
+    # Print each hyperparam
+    for key, val in task_cfg.items():
+        if key in skip:
+            continue
+
+        # Special PGD block header
+        if task == "adversarial_attack" and key == "attack":
+            print(f"\n{RED}{BOLD}PGD attack hyper-parameters{RESET}")
+
+        # Nested dicts
+        if isinstance(val, (DictConfig, dict)):
+            print(f"{BLUE}{BOLD}{key}{RESET}:")
+            for subkey, subval in val.items():
+                print(f"  {WHITE}{subkey}{RESET}: {subval}")
+        else:
+            print(f"{BLUE}{key}{RESET}: {WHITE}{val}{RESET}")
+
+    print(sep)
+
+
 
 def get_hyperaparams_dict(cfg: DictConfig) -> dict:
     """