Walkthrough: one `train` run, step by step¶

This page follows one run of python scripts/launcher.py train research/stable/ordinal_regression_mtl_mar_2024/config.yaml (DAM 1, audio only) in execution order. Steps 1-20 run once at start-up (13-20 once per split: train, val, test): read the config, build the model, loaders and datasets. Inside trainer.fit, steps 21-23 run per recording and 24-35 per batch: audio windows, Whisper encoder, head, loss. Steps 36-42 run after each validation pass: tune thresholds, compute metrics, save the checkpoint. trainer.test then runs the same path on the test split with the val-tuned thresholds (no tuning); step 43 runs at the end and step 44 is a separate analysis call. As shipped, the config stops at steps 2, 11 and 22; each "Watch out" gives the fix. Each step shows where the function is, what it does (its docstring) and what it receives and returns in this run (B = recordings per batch = 2, W = windows per recording, ΣW = windows in the batch). The numbered gold boxes in the Code map are these same steps.

flowchart TB subgraph t_launch["1 · Launch"] direction LR s1["1 · launcher()"] class s1 main click s1 "#step-1" s2["2 · _setup_launcher()"] class s2 main click s2 "#step-2" s3["3 · LauncherConfig.__init__()"] class s3 main click s3 "#step-3" s4["4 · TrainingParams()"] class s4 main click s4 "#step-4" s5["5 · LauncherConfig._init_data_params()"] class s5 main click s5 "#step-5" s6["6 · LauncherConfig.get_params_by_split()"] class s6 main click s6 "#step-6" s7["7 · LauncherConfig._init_model()"] class s7 main click s7 "#step-7" s8["8 · DepAnxClassifier()"] class s8 main click s8 "#step-8" s9["9 · KintsugiBackboneBase.apply_lora()"] class s9 main click s9 "#step-9" s10["10 · _setup_callbacks()"] class s10 main click s10 "#step-10" s11["11 · _log_to_wandb()"] class s11 main click s11 "#step-11" s12["12 · run_training()"] class s12 main click s12 "#step-12" s13["13 · LauncherConfig.get_data_loader()"] class s13 main click s13 "#step-13" s14["14 · LauncherConfig.get_dataset()"] class s14 main click s14 "#step-14" end subgraph t_data["2 · Data"] direction LR s15["15 · KintsugiMetadata()"] class s15 main click s15 "#step-15" s16["16 · load_relative_dataset()"] class s16 main click s16 "#step-16" s17["17 · KintsugiMetadata.preprocess_metadata()"] class s17 main click s17 "#step-17" s18["18 · KintsugiMetadata.filter_metadata()"] class s18 main click s18 "#step-18" s19["19 · get_preprocessor_with_audio_normalization()"] class s19 main click s19 "#step-19" s20["20 · UuidSampler()"] class s20 main click s20 "#step-20" s21["21 · KintsugiMetadata.__getitem__()"] class s21 main click s21 "#step-21" s22["22 · KintsugiAudioDataset.__getitem__()"] class s22 main click s22 "#step-22" s23["23 · get_window_starts()"] class s23 main click s23 "#step-23" s24["24 · broadcast_collate()"] class s24 main click s24 "#step-24" end subgraph t_model["3 · Model (DAM 1)"] direction LR s25["25 · MultiTaskOrdinalRegressionPLModule.training_step()"] class s25 main click s25 "#step-25" s26["26 · MultiTaskOrdinalRegressionPLModule.step()"] class s26 main click s26 "#step-26" s27["27 · DepAnxClassifier.forward()"] class s27 main click s27 "#step-27" s28["28 · DepAnxClassifier.forward_backbone()"] class s28 main click s28 "#step-28" s29["29 · WhisperTorchFeatureExtractor.forward()"] class s29 main click s29 "#step-29" s30["30 · WhisperBackbone.forward()"] class s30 main click s30 "#step-30" s31["31 · MultitaskHead.forward()"] class s31 main click s31 "#step-31" end subgraph t_loss["4 · Loss"] direction LR s32["32 · OrdinalRegressionPLModule.step()"] class s32 main click s32 "#step-32" s33["33 · OrdinalRegressionPLModule.compute_and_log_loss()"] class s33 main click s33 "#step-33" s34["34 · average_and_variance_tensor_in_segments()"] class s34 main click s34 "#step-34" s35["35 · CORALLoss.forward()"] class s35 main click s35 "#step-35" end subgraph t_metrics["5 · Thresholds and metrics"] direction LR s36["36 · MultiTaskOrdinalRegressionPLModule.on_validation_test_epoch_end()"] class s36 main click s36 "#step-36" s37["37 · OrdinalRegressionPLModule.on_validation_test_epoch_end()"] class s37 main click s37 "#step-37" s38["38 · OrdinalRegressionPLModule.quantize_labels()"] class s38 main click s38 "#step-38" s39["39 · IndetSnSpArray.build()"] class s39 main click s39 "#step-39" s40["40 · OptimalOrdinalThresholdingViaDynamicProgramming.tune_thresholds()"] class s40 main click s40 "#step-40" s41["41 · OrdinalRegressionPLModule.log_ordinal_metrics()"] class s41 main click s41 "#step-41" end subgraph t_outputs["6 · Outputs and analysis"] direction LR s42["42 · DistributedModelCheckpoint.on_validation_end()"] class s42 main click s42 "#step-42" s43["43 · _log_datasets_to_wandb()"] class s43 main click s43 "#step-43" s44["44 · get_dataframe_with_scores()"] class s44 main click s44 "#step-44" end s1 ==> s2 s2 ==> s3 s3 ==> s4 s4 ==> s5 s5 ==> s6 s6 ==> s7 s7 ==> s8 s8 ==> s9 s9 ==> s10 s10 ==> s11 s11 ==> s12 s12 ==> s13 s13 ==> s14 s14 ==> s15 s15 ==> s16 s16 ==> s17 s17 ==> s18 s18 ==> s19 s19 ==> s20 s20 ==> s21 s21 ==> s22 s22 ==> s23 s23 ==> s24 s24 ==> s25 s25 ==> s26 s26 ==> s27 s27 ==> s28 s28 ==> s29 s29 ==> s30 s30 ==> s31 s31 ==> s32 s32 ==> s33 s33 ==> s34 s34 ==> s35 s35 ==> s36 s36 ==> s37 s37 ==> s38 s38 ==> s39 s39 ==> s40 s40 ==> s41 s41 ==> s42 s42 ==> s43 s43 ==> s44 classDef main fill:#FFF4DC,stroke:#A67C1B,stroke-width:2px,color:#1A1A1A classDef opt fill:#F3F4F6,stroke:#B8BEC7,color:#555,stroke-dasharray:4 3 classDef edge fill:#FFFFFF,stroke:#1F7A5C,color:#1F7A5C,stroke-width:1.5px

Step 1 · `launcher`¶

scripts/launcher.py:43 · open code ↗ · stage 1 · Launch

Does: Run one mode (train/evaluate/cache/compile) end to end from a YAML config.

In	Out
mode "train", config path .../ordinal_regression_mtl_mar_2024/config.yaml	calls, in order: _setup_launcher (step 2), _setup_callbacks (10), _log_to_wandb (11), then builds Trainer (GPU fp16 or CPU fp32, DDP, max_epochs 32, grad accumulation) and calls run_training (12); after it returns, _log_datasets_to_wandb (43)

Watch out

Validation runs every 16000 // (2 × GPUs) training batches, counted across epochs (check_val_every_n_epoch None), not once per epoch; 16000 must divide by the global batch size. If all 32 epochs together have fewer batches than that, validation never runs, no best checkpoint exists and run_training crashes at torch.load. Lower training_examples_per_eval for small datasets.

Step 2 · _setup_launcher →

Step 2 · `_setup_launcher`¶

scripts/launcher.py:155 · open code ↗ · stage 1 · Launch

Does: Load the YAML into LauncherConfig and create the W&B logger on rank 0 only.

In	Out
from step 1: config path, experiment_dir None	LauncherParams(launcher_config, launcher_logger); building LauncherConfig runs steps 3-9; the W&B logger exists only on rank 0

Watch out

wandb_entity and wandb_project are hardcoded to None (launcher.py) and raise ValueError right after LauncherConfig (model already built, Whisper already downloaded): set them before the first run.

← Step 1 · launcher | Step 3 · LauncherConfig.__init__ →

Step 3 · `LauncherConfig.init`¶

kirad/launcher_utils.py:234 · open code ↗ · stage 1 · Launch

Does: Read YAML, parse training/data/model params and instantiate the model.

In	Out
from step 2: config path; experiment_dir None -> the config's folder	LauncherConfig with training_params, model_params, model; experiment folder added to sys.path

Watch out

model_path "model.DepAnxClassifier" is imported from the experiment folder via sys.path, so it resolves to that folder's model.py.

← Step 2 · _setup_launcher | Step 4 · TrainingParams →

Step 4 · `TrainingParams`¶

kirad/launcher_utils.py:101 · open code ↗ · stage 1 · Launch

Does: Validated training_params: batch sizes, epochs, monitor metric, optimizer, seed.

In	Out
from step 3: training_params: batch_size 2, effective_batch_size 128, max_epochs 32, seed 44, optimizer {backbone lr 1.6e-4, classifier lr 8e-4, wd 1e-3}, scheduler, tasks	validated params; _global_batch_size = WORLD_SIZE × 2; pl.seed_everything(44)

Watch out

Grad accumulation = 128 // (GPUs × 2), e.g. 64 on 1 GPU; must divide exactly or the launcher raises.

← Step 3 · LauncherConfig.__init__ | Step 5 · LauncherConfig._init_data_params →

Step 5 · `LauncherConfig._init_data_params`¶

kirad/launcher_utils.py:331 · open code ↗ · stage 1 · Launch

Does: Parse data_params keys "<modality>=<class.path>" into factories + per-split params.

In	Out
from step 3: data_params keys "metadata=...KintsugiMetadata" and "audio=...KintsugiAudioDataset"	metadata_factory, {"audio": KintsugiAudioDataset}, params per split (train/val/test)

← Step 4 · TrainingParams | Step 6 · LauncherConfig.get_params_by_split →

Step 6 · `LauncherConfig.get_params_by_split`¶

kirad/launcher_utils.py:533 · open code ↗ · stage 1 · Launch

Does: Merge config params for one split: default < val_test (val/test) < split.

In	Out
from step 5: the whole mapping of one data_params key (sections default, val_test, train), split "train", "val" or "test"	merged dict, e.g. metadata train: duration_column "prompt_vad_trimmed_duration", merge_rows_by_subject "many_to_one", dataset_splits "train"; audio train: window_method "random_adjacent_pair"

Watch out

Any string value that is an existing path (file or folder, after ~ expansion; relative ones against the experiment folder) is replaced by its absolute Path.

← Step 5 · LauncherConfig._init_data_params | Step 7 · LauncherConfig._init_model →

Step 7 · `LauncherConfig._init_model`¶

kirad/launcher_utils.py:372 · open code ↗ · stage 1 · Launch

Does: Import model_path, build the model and load ckpt weights (or defer full restore).

In	Out
from step 3: model_params (model_path "model.DepAnxClassifier", ckpt_path "" -> None), config (whisper_config, classifier_config), training_params	self.model = DepAnxClassifier; no weights loaded (ckpt_path None); restore_training_state_ckpt_path None

← Step 6 · LauncherConfig.get_params_by_split | Step 8 · DepAnxClassifier →

Step 8 · `DepAnxClassifier`¶

research/stable/ordinal_regression_mtl_mar_2024/model.py:25 · open code ↗ · stage 1 · Launch

Does: DAM 1: log-mel -> Whisper small.en encoder (LoRA) -> mean pool -> multitask CORAL head.

In	Out
from step 7 (import_object): whisper_config (small.en, LoRA), classifier_config [256, 64] + tasks, training_params.tasks	LightningModule: tasks {depression: 28 classes, anxiety: 22}, feature_extractor, backbone, head

Watch out

Whisper small.en is downloaded from HuggingFace here. No featex_config in the YAML, so the GPU extractor keeps normalize_audio and ideal_logmel_energies on.

← Step 7 · LauncherConfig._init_model | Step 9 · KintsugiBackboneBase.apply_lora →

Step 9 · `KintsugiBackboneBase.apply_lora`¶

kirad/base_models.py:818 · open code ↗ · stage 1 · Launch

Does: Wrap self.backbone with PEFT LoRA from lora_params; no-op (full fine-tune) if empty.

In	Out
from step 8 (WhisperBackbone.init): lora_params: r 32, lora_alpha 64.0, "all-linear", lora_dropout 0.4, modules_to_save [conv1, conv2], bias "all"	self.backbone = PEFT-wrapped Whisper encoder; logs the LoRA-targeted and fully trained modules

Watch out

Called from WhisperBackbone.__init__. bias "all" also trains every bias of the encoder; conv1/conv2 are trained in full.

← Step 8 · DepAnxClassifier | Step 10 · _setup_callbacks →

Step 10 · `_setup_callbacks`¶

scripts/launcher.py:191 · open code ↗ · stage 1 · Launch

Does: Best-checkpoint saver, early stopping and LR monitor on "mode:metric" monitor_loss.

In	Out
from step 1 (launcher): monitor_loss "max:metrics/val/sn_eq_sp@0%i", early_stopping_patience 10	[LearningRateMonitor(step), EarlyStopping(max, patience 10), DistributedModelCheckpoint(top-1 + last)]

Watch out

Patience counts validations (every 16000 training examples), not epochs.

← Step 9 · KintsugiBackboneBase.apply_lora | Step 11 · _log_to_wandb →

Step 11 · `_log_to_wandb`¶

scripts/launcher.py:232 · open code ↗ · stage 1 · Launch

Does: Upload experiment + kirad source, package versions and ideal log-mel energies to W&B.

In	Out
from step 1: LauncherConfig, LauncherLogger from step 2 (rank 0 only)	W&B artifacts source-experiment-<run>, source-kirad-<run>, package_info table

Watch out

get_package_info("kipy") fails without the private kipy package; remove "kipy" from the list in launcher.py.

← Step 10 · _setup_callbacks | Step 12 · run_training →

Step 12 · `run_training`¶

scripts/launcher.py:119 · open code ↗ · stage 1 · Launch

Does: Fit on train/val, reload the best checkpoint, then test it on the test split.

In	Out
from step 1: Trainer, LauncherConfig	trainer.fit(model, train, val) builds loaders via steps 13-20 and runs 21-42; reload best_model_path weights; trainer.test on the test loader (same path, split test, val-tuned thresholds)

Watch out

Inside fit, Lightning calls configure_model (adds the <task>/ log prefix), configure_optimizers (AdamW: backbone lr 1.6e-4, head + CORAL biases 8e-4; MultiStepLR halves both at epochs 20 and 26) and runs 2 sanity val batches.

← Step 11 · _log_to_wandb | Step 13 · LauncherConfig.get_data_loader →

Step 13 · `LauncherConfig.get_data_loader`¶

kirad/launcher_utils.py:464 · open code ↗ · stage 1 · Launch

Does: DataLoader for a split (cached per split and mode); batch_size = recordings.

In	Out
from step 12: split "train", training_mode True (then "val" and "test" with False)	DataLoader: batch_size 2, broadcast_collate, UuidSampler; drop_last True for train

Watch out

num_workers = (CPUs - 1) // GPUs unless training_params.cpus_per_worker is set.

← Step 12 · run_training | Step 14 · LauncherConfig.get_dataset →

Step 14 · `LauncherConfig.get_dataset`¶

kirad/launcher_utils.py:418 · open code ↗ · stage 1 · Launch

Does: StackDataset of metadata + each modality dataset for a split (built once).

In	Out
from step 13: split "train"	StackDataset(metadata=KintsugiMetadata, audio=KintsugiAudioDataset); each item is a dict with keys "metadata" and "audio"

Watch out

Items are indexed by UUID string (from UuidSampler), not by integer position.

← Step 13 · LauncherConfig.get_data_loader | Step 15 · KintsugiMetadata →

Step 15 · `KintsugiMetadata`¶

kirad/datasets_and_dataloaders.py:71 · open code ↗ · stage 2 · Data

Does: Metadata table indexed by UUID: loads CSV, merges KD scores, filters by duration.

In	Out
from step 14: train: "v2.11.4", 4 data_sources, fold 0, label phq/gad, prompt_vad_trimmed_duration, many_to_one	metadata object; _df indexed by uuid, preprocessed and filtered

Watch out

"v2.11.4" is a private Kintsugi dataset version fetched with gsutil; for new data set df_or_dataset_version to a CSV path.

← Step 14 · LauncherConfig.get_dataset | Step 16 · load_relative_dataset →

Step 16 · `load_relative_dataset`¶

kirad/dataset_utils.py:97 · open code ↗ · stage 2 · Data

Does: Load metadata and add audio_filename / text_filename paths from the filename column.

In	Out
from step 15 (KintsugiMetadata.init): "v2.11.4", data_sources, "train", fold 0, audio_root ~/.kintsugi/audio/resampled_and_trimmed_...	DataFrame + audio_filename <audio_root>/<filename>.wav, text_filename <cwd>/<filename>.json

← Step 15 · KintsugiMetadata | Step 17 · KintsugiMetadata.preprocess_metadata →

Step 17 · `KintsugiMetadata.preprocess_metadata`¶

kirad/datasets_and_dataloaders.py:210 · open code ↗ · stage 2 · Data

Does: Join KD score targets, optionally merge rows per subject, count windows per recording.

In	Out
from step 16: _df; train: merge_rows_by_subject many_to_one; val/test: False	train: one row per subject (durations summed, filename list); all splits: + num_windows column

Watch out

merge_df_rows_by_subject drops audio_filename, so with this config the train KintsugiAudioDataset.__getitem__ (step 22) fails on row.audio_filename.

← Step 16 · load_relative_dataset | Step 18 · KintsugiMetadata.filter_metadata →

Step 18 · `KintsugiMetadata.filter_metadata`¶

kirad/datasets_and_dataloaders.py:257 · open code ↗ · stage 2 · Data

Does: Subsample fraction of rows, then drop recordings outside min/max duration (s).

In	Out
from step 17: _df, fraction 1.0, min_duration 30 s (default), max_duration None	_df sampled (random_state 1) and limited to recordings with duration_column >= 30 s

← Step 17 · KintsugiMetadata.preprocess_metadata | Step 19 · get_preprocessor_with_audio_normalization →

Step 19 · `get_preprocessor_with_audio_normalization`¶

kirad/datasets_and_dataloaders.py:790 · open code ↗ · stage 2 · Data

Does: Build the function audio [T] -> windows -> features used by KintsugiAudioDataset.

In	Out
from step 14 (KintsugiAudioDataset.init): dummy_preprocessor (preprocessor false), window_method, normalize_audio True, 30 s, max_overlap_frac 0.0	closure (audio [T], sr) -> raw windows [W, 480000], DC removed and peak-scaled to [-1, 1]

Watch out

Built once in KintsugiAudioDataset.__init__ (init_data_processor), which step 14 triggers; the closure runs on every item.

← Step 18 · KintsugiMetadata.filter_metadata | Step 20 · UuidSampler →

Step 20 · `UuidSampler`¶

kirad/datasets_and_dataloaders.py:703 · open code ↗ · stage 2 · Data

Does: Yields metadata UUIDs (shuffled for train, sequential for val/test) for the loader.

In	Out
from step 13: KintsugiMetadata; shuffle True (train) or False (val/test)	iterator of UUID strings (RandomSampler without replacement, or SequentialSampler)

← Step 19 · get_preprocessor_with_audio_normalization | Step 21 · KintsugiMetadata.__getitem__ →

Step 21 · `KintsugiMetadata.getitem`¶

kirad/datasets_and_dataloaders.py:293 · open code ↗ · stage 2 · Data

Does: Return {uuid, label: {task: {primary: raw score, kd: teacher score}}} for one UUID.

In	Out
uuid (str) from UuidSampler (step 20). From here on the code runs inside trainer.fit (step 12): steps 21-23 once per recording, steps 24-35 once per batch	{uuid, label: {depression: {primary: phq}, anxiety: {primary: gad}}}; no kd key in this run

Watch out

Labels are raw PHQ-9 (0-27) and GAD-7 (0-21) scores; binning happens only in the metrics.

← Step 20 · UuidSampler | Step 22 · KintsugiAudioDataset.__getitem__ →

Step 22 · `KintsugiAudioDataset.getitem`¶

kirad/datasets_and_dataloaders.py:439 · open code ↗ · stage 2 · Data

Does: Load, concatenate (channel 0), augment and window one recording's audio.

In	Out
uuid (str), same as step 21	raw audio windows: train [2, 480000] ("random_adjacent_pair", once many_to_one is off); val/test [W, 480000] ("all")

Watch out

Train split with this config fails here with AttributeError on row.audio_filename: many_to_one merging (step 17) dropped it; set merge_rows_by_subject False to run it. Not resampled: a sample rate other than 16 kHz only logs a warning, windows become sr × 30 samples and forward_backbone (step 28) silently crops or pads them to 480000. Only channel 0 is kept. No augmentations configured.

← Step 21 · KintsugiMetadata.__getitem__ | Step 23 · get_window_starts →

Step 23 · `get_window_starts`¶

kirad/datasets_and_dataloaders.py:738 · open code ↗ · stage 2 · Data

Does: Start sample(s) of the windows cut from one recording, by window_method.

In	Out
from step 19 (the closure in step 22): method, audio_len in samples, 480000, max_overlap_frac 0.0	train: [s, s + 480000] with s random if audio >= 60 s, else fixed [0, len - 480000] (overlapping); val/test: 1 + (len - 480000) // 480000 starts spread evenly (np.linspace) from 0 to len - 480000

← Step 22 · KintsugiAudioDataset.__getitem__ | Step 24 · broadcast_collate →

Step 24 · `broadcast_collate`¶

kirad/dataset_utils.py:830 · open code ↗ · stage 2 · Data

Does: DataLoader collate: concat windows of all recordings; length = windows per recording.

In	Out
list of 2 items {"metadata": {uuid, label}, "audio": [W_i, 480000]}	{uuid [2], label.<task>.primary [2], length [2], features.audio [ΣW, 480000]}; train ΣW = 4

Watch out

Windows are concatenated, not stacked; length is the only link from windows back to recordings.

← Step 23 · get_window_starts | Step 25 · MultiTaskOrdinalRegressionPLModule.training_step →

Step 25 · `MultiTaskOrdinalRegressionPLModule.training_step`¶

kirad/base_models.py:722 · open code ↗ · stage 3 · Model (DAM 1)

Does: Lightning train step: returns the step dict with total_loss.

In	Out
from step 24: batch dict, batch_idx	step("train") dict with loss = total_loss; Lightning backprops; AdamW steps every accumulate_grad_batches batches

Watch out

Lightning calls training_step, validation_step and test_step, not your code; validation_step runs the same path with split val.

← Step 24 · broadcast_collate | Step 26 · MultiTaskOrdinalRegressionPLModule.step →

Step 26 · `MultiTaskOrdinalRegressionPLModule.step`¶

kirad/base_models.py:734 · open code ↗ · stage 3 · Model (DAM 1)

Does: Run model once, sum each task's loss (task.step) into total_loss.

In	Out
split "train", batch dict from step 24 with uuid, features.audio [ΣW, 480000], label, length [B]	total_loss = sum of task.step losses (step 32); logs metrics/train/total_loss; returns a dict with uuid, features, label, length, loss

← Step 25 · MultiTaskOrdinalRegressionPLModule.training_step | Step 27 · DepAnxClassifier.forward →

Step 27 · `DepAnxClassifier.forward`¶

research/stable/ordinal_regression_mtl_mar_2024/model.py:84 · open code ↗ · stage 3 · Model (DAM 1)

Does: Backbone then MultitaskHead -> ({task: [ΣW, 1] scores}, lengths).

In	Out
from step 26: x {"audio": [ΣW, 480000]}, lengths [B]	({depression: [ΣW, 1], anxiety: [ΣW, 1]}, lengths)

← Step 26 · MultiTaskOrdinalRegressionPLModule.step | Step 28 · DepAnxClassifier.forward_backbone →

Step 28 · `DepAnxClassifier.forward_backbone`¶

research/stable/ordinal_regression_mtl_mar_2024/model.py:89 · open code ↗ · stage 3 · Model (DAM 1)

Does: Pad raw audio to 480,000 samples, log-mel on GPU, Whisper encoder -> [ΣW, 768].

In	Out
x["audio"] [ΣW, 480000], cached_backbone False	[ΣW, 768] one embedding per window; calls step 29 (log-mel), then step 30 (encoder)

← Step 27 · DepAnxClassifier.forward | Step 29 · WhisperTorchFeatureExtractor.forward →

Step 29 · `WhisperTorchFeatureExtractor.forward`¶

kirad/feature_extraction_whisper.py:138 · open code ↗ · stage 3 · Model (DAM 1)

Does: Normalize each window, compute log-mel, shift bins to IDEAL_LOGMEL_ENERGIES.

In	Out
from step 28: [ΣW, 480000] padded raw audio, float32 (fp16 autocast does not cast inputs)	[ΣW, 80, 3000] log-mel, float32 (the input dtype)

Watch out

Computed in float64; each window is normalized to zero mean, unit variance, then shifted to the fixed IDEAL_LOGMEL_ENERGIES profile.

← Step 28 · DepAnxClassifier.forward_backbone | Step 30 · WhisperBackbone.forward →

Step 30 · `WhisperBackbone.forward`¶

kirad/base_models.py:877 · open code ↗ · stage 3 · Model (DAM 1)

Does: Whisper encoder on log-mel windows, mean over time -> [ΣW, D] (768 for small.en).

In	Out
from step 29:	last_hidden_state [ΣW, 1500, 768] -> mean over time -> [ΣW, 768]

Watch out

LoRA (r 32, alpha 64) sits in this encoder; dropout, activation_dropout and encoder_layerdrop are 0.0 from hf_config.

← Step 29 · WhisperTorchFeatureExtractor.forward | Step 31 · MultitaskHead.forward →

Step 31 · `MultitaskHead.forward`¶

kirad/base_models.py:1075 · open code ↗ · stage 3 · Model (DAM 1)

Does: Shared layers, then one TaskHead per task -> {task: [ΣW, 1] score}.

In	Out
from step 28: [ΣW, 768]	{depression, anxiety}: [ΣW, 1]; 768 -> 256 Mish -> 64, then 64 -> 128 Mish Dropout 0.4 -> 1 (no bias)

← Step 30 · WhisperBackbone.forward | Step 32 · OrdinalRegressionPLModule.step →

Step 32 · `OrdinalRegressionPLModule.step`¶

kirad/base_models.py:572 · open code ↗ · stage 4 · Loss

Does: Accumulate per-recording mean/var scores for val/test, then return the loss.

In	Out
from step 26, once per task: split, uuid, features = this task's [ΣW, 1] scores from step 27, label primary [B], length [B]	task loss (scalar); on val/test also accumulates uuid, mean score [B], variance [B], label [B]

Watch out

This module's forward is the identity: the scores already come from the shared head.

← Step 31 · MultitaskHead.forward | Step 33 · OrdinalRegressionPLModule.compute_and_log_loss →

Step 33 · `OrdinalRegressionPLModule.compute_and_log_loss`¶

kirad/base_models.py:504 · open code ↗ · stage 4 · Loss

Does: Loss = coral_w·CORAL + var_w·Var(window scores) + kd_w·MSE(teacher), logged.

In	Out
from step 32: logits [ΣW, 1], y {primary [B]}, lengths [B], split "train"	1.0 × CORAL + 40.0 × mean window-score variance; logs <task>/metrics/train/loss_exact28 (exact22 for anxiety)

Watch out

Train: CORAL per window against the recording's label, then averaged. Val/test: CORAL of the per-recording mean score. No KD term (no score_targets_csv).

← Step 32 · OrdinalRegressionPLModule.step | Step 34 · average_and_variance_tensor_in_segments →

Step 34 · `average_and_variance_tensor_in_segments`¶

kirad/utils.py:122 · open code ↗ · stage 4 · Loss

Does: Per-recording mean and (biased) variance of window scores: [ΣW, ...] -> two [B, ...].

In	Out
from step 33: logits [ΣW, 1], lengths [B]	mean [B, 1] and biased variance [B, 1] of the window scores per recording

← Step 33 · OrdinalRegressionPLModule.compute_and_log_loss | Step 35 · CORALLoss.forward →

Step 35 · `CORALLoss.forward`¶

kirad/losses.py:43 · open code ↗ · stage 4 · Loss

Does: Weighted sum of BCE over cutoffs: sigmoid(s + b_k) vs [y > k] for each cutoff k.

In	Out
from step 33: x [ΣW, 1] (train) or [B, 1] (val/test), y label [N]	[N] loss = sum over k of BCE(s + b_k, y > k); 27 biases (depression), 21 (anxiety), all weights 1

← Step 34 · average_and_variance_tensor_in_segments | Step 36 · MultiTaskOrdinalRegressionPLModule.on_validation_test_epoch_end →

Step 36 · `MultiTaskOrdinalRegressionPLModule.on_validation_test_epoch_end`¶

kirad/base_models.py:680 · open code ↗ · stage 5 · Thresholds and metrics

Does: Run each task's epoch end, then log harmonic mean of shared metrics across tasks.

In	Out
split "val", after Lightning has run all val batches through steps 26-35 (validation_step; step 32 accumulates scores) every val_check_interval batches (step 1)	runs steps 37-41 for each task, then logs the harmonic mean over tasks of each shared metric as metrics/val/..., e.g. metrics/val/sn_eq_sp@0%i

Watch out

That combined metric picks the best checkpoint and drives early stopping.

← Step 35 · CORALLoss.forward | Step 37 · OrdinalRegressionPLModule.on_validation_test_epoch_end →

Step 37 · `OrdinalRegressionPLModule.on_validation_test_epoch_end`¶

kirad/base_models.py:423 · open code ↗ · stage 5 · Thresholds and metrics

Does: Gather scores, log W&B table, tune thresholds on val, compute Sn/Sp metrics.

In	Out
split "val"; accumulated uuids, mean scores, variances, labels [N_val] (deduplicated)	W&B artifact scores-<task>-val-<run_id>; tuned thresholds and Sn=Sp buffers; metrics dict

Watch out

Lightning's sanity check (2 val batches before training) also runs this and tunes thresholds; only the W&B table is skipped.

← Step 36 · MultiTaskOrdinalRegressionPLModule.on_validation_test_epoch_end | Step 38 · OrdinalRegressionPLModule.quantize_labels →

Step 38 · `OrdinalRegressionPLModule.quantize_labels`¶

kirad/base_models.py:222 · open code ↗ · stage 5 · Thresholds and metrics

Does: Map input labels (e.g. PHQ-9 sum) to output labels (e.g. 0/1/2 for L/M/H).

In	Out
labels [N_val] raw PHQ-9 / GAD-7	classes [N_val]: depression 0-2 (> 9, > 14), anxiety 0-3 (> 4, > 9, > 14)

← Step 37 · OrdinalRegressionPLModule.on_validation_test_epoch_end | Step 39 · IndetSnSpArray.build →

Step 39 · `IndetSnSpArray.build`¶

kirad/metrics.py:176 · open code ↗ · stage 5 · Thresholds and metrics

Does: Sn/Sp/indet_frac for all threshold pairs (or the given ones) via histogram cumsums.

In	Out
thresholds None (<= 10,000 scores), y_score [N_val] float64, y_true = label > cut	Sn, Sp, indet_frac for every (lower, upper) pair of score midpoints; one array per binary cut

Watch out

Pairs grow as N_val squared; above 10,000 scores a 10,000-point linspace grid is used instead.

← Step 38 · OrdinalRegressionPLModule.quantize_labels | Step 40 · OptimalOrdinalThresholdingViaDynamicProgramming.tune_thresholds →

Step 40 · `OptimalOrdinalThresholdingViaDynamicProgramming.tune_thresholds`¶

kirad/ordinal_thresholding.py:179 · open code ↗ · stage 5 · Thresholds and metrics

Does: Set thresholds by DP over midpoints of unique scores (+/-inf); return the best metric.

In	Out
scores [N_val], quantized labels [N_val]	self.thresholds (K - 1 values) that optimize the method's metric; returns that metric's optimal mean cost

Watch out

Tuned on val only (5 DP methods; CORAL copies -b_k at cuts 9, 14 / 4, 9, 14). Stored as buffers in the checkpoint and reused on test.

← Step 39 · IndetSnSpArray.build | Step 41 · OrdinalRegressionPLModule.log_ordinal_metrics →

Step 41 · `OrdinalRegressionPLModule.log_ordinal_metrics`¶

kirad/base_models.py:228 · open code ↗ · stage 5 · Thresholds and metrics

Does: Log AUROC and min(Sn, Sp) per binary cut and indet budget, plus their mean.

In	Out
split "val", one IndetSnSpArray per cut, scores, labels	logs <task>/bucketed_metrics/val/auroc_>9@0%i, sn_eq_sp_>9@40%i, ... and <task>/metrics/val/sn_eq_sp@0%i

Watch out

On test it also logs Sn, Sp and indet fraction at the thresholds tuned on val. Right after it, log_thresholding_metrics logs accuracy / F1 / recall etc. at the thresholds of each of the six tuning methods.

← Step 40 · OptimalOrdinalThresholdingViaDynamicProgramming.tune_thresholds | Step 42 · DistributedModelCheckpoint.on_validation_end →

Step 42 · `DistributedModelCheckpoint.on_validation_end`¶

kirad/launcher_utils.py:45 · open code ↗ · stage 6 · Outputs and analysis

Does: Save best/last ckpt after validation (with tuned thresholds); replicate on nodes.

In	Out
trainer, model; metric metrics/val/sn_eq_sp@0%i (mode max). Lightning calls this after each validation	writes chkpts/dam1_audio_concat_overlap_0.0_variance_weight_40.0/epoch-E_metric-0.000.ckpt (best) + last.ckpt

Watch out

Runs after the val epoch end, so the checkpoint holds the val-tuned thresholds. The filename is a plain string, so {monitor_loss_name} is not a logged metric and Lightning writes 0.000: the name never shows the score. Uses torch.distributed barriers; /mnt/chkpts if JOB_NAME is set.

← Step 41 · OrdinalRegressionPLModule.log_ordinal_metrics | Step 43 · _log_datasets_to_wandb →

Step 43 · `_log_datasets_to_wandb`¶

scripts/launcher.py:286 · open code ↗ · stage 6 · Outputs and analysis

Does: Upload the metadata table actually used for each split as a W&B "dataset" artifact.

In	Out
back in launcher (step 1) after run_training returns: LauncherConfig (train/val/test datasets built), LauncherLogger	W&B artifact dataset-<run_id> with train_dataset, val_dataset, test_dataset metadata tables

← Step 42 · DistributedModelCheckpoint.on_validation_end | Step 44 · get_dataframe_with_scores →

Step 44 · `get_dataframe_with_scores`¶

kirad/utils.py:442 · open code ↗ · stage 6 · Outputs and analysis

Does: Join a run's metadata CSVs with its W&B scores-<task>-<split> tables on uuid.

In	Out
wandb_path "entity/project/run_id", split "val" or "test"	metadata DataFrame joined on uuid with scores_<task> and quantized labels from the W&B tables

Watch out

After the run, not called by the launcher. Needs the model-<run_id> artifact (tag best) from WandbLogger(log_model=True). For val it reads scores version v{checkpoint epoch}, only right if validation runs once per epoch; with training_examples_per_eval (this config) pass scores_tag. Test reads v0.

← Step 43 · _log_datasets_to_wandb

Walkthrough: one train run, step by step¶

Step 1 · launcher¶

Step 2 · _setup_launcher¶

Step 3 · LauncherConfig.__init__¶

Step 4 · TrainingParams¶

Step 5 · LauncherConfig._init_data_params¶

Step 6 · LauncherConfig.get_params_by_split¶

Step 7 · LauncherConfig._init_model¶

Step 8 · DepAnxClassifier¶

Step 9 · KintsugiBackboneBase.apply_lora¶

Step 10 · _setup_callbacks¶

Step 11 · _log_to_wandb¶

Step 12 · run_training¶

Step 13 · LauncherConfig.get_data_loader¶

Step 14 · LauncherConfig.get_dataset¶

Step 15 · KintsugiMetadata¶

Step 16 · load_relative_dataset¶

Step 17 · KintsugiMetadata.preprocess_metadata¶

Step 18 · KintsugiMetadata.filter_metadata¶

Step 19 · get_preprocessor_with_audio_normalization¶

Step 20 · UuidSampler¶

Step 21 · KintsugiMetadata.__getitem__¶

Step 22 · KintsugiAudioDataset.__getitem__¶

Step 23 · get_window_starts¶

Step 24 · broadcast_collate¶

Step 25 · MultiTaskOrdinalRegressionPLModule.training_step¶

Step 26 · MultiTaskOrdinalRegressionPLModule.step¶

Step 27 · DepAnxClassifier.forward¶

Step 28 · DepAnxClassifier.forward_backbone¶

Step 29 · WhisperTorchFeatureExtractor.forward¶

Step 30 · WhisperBackbone.forward¶

Step 31 · MultitaskHead.forward¶

Step 32 · OrdinalRegressionPLModule.step¶

Step 33 · OrdinalRegressionPLModule.compute_and_log_loss¶

Step 34 · average_and_variance_tensor_in_segments¶

Step 35 · CORALLoss.forward¶

Step 36 · MultiTaskOrdinalRegressionPLModule.on_validation_test_epoch_end¶

Step 37 · OrdinalRegressionPLModule.on_validation_test_epoch_end¶

Step 38 · OrdinalRegressionPLModule.quantize_labels¶

Step 39 · IndetSnSpArray.build¶

Step 40 · OptimalOrdinalThresholdingViaDynamicProgramming.tune_thresholds¶

Step 41 · OrdinalRegressionPLModule.log_ordinal_metrics¶

Step 42 · DistributedModelCheckpoint.on_validation_end¶

Step 43 · _log_datasets_to_wandb¶

Step 44 · get_dataframe_with_scores¶

Walkthrough: one `train` run, step by step¶

Step 1 · `launcher`¶

Step 2 · `_setup_launcher`¶

Step 3 · `LauncherConfig.init`¶

Step 4 · `TrainingParams`¶

Step 5 · `LauncherConfig._init_data_params`¶

Step 6 · `LauncherConfig.get_params_by_split`¶

Step 7 · `LauncherConfig._init_model`¶

Step 8 · `DepAnxClassifier`¶

Step 9 · `KintsugiBackboneBase.apply_lora`¶

Step 10 · `_setup_callbacks`¶

Step 11 · `_log_to_wandb`¶

Step 12 · `run_training`¶

Step 13 · `LauncherConfig.get_data_loader`¶

Step 14 · `LauncherConfig.get_dataset`¶

Step 15 · `KintsugiMetadata`¶

Step 16 · `load_relative_dataset`¶

Step 17 · `KintsugiMetadata.preprocess_metadata`¶

Step 18 · `KintsugiMetadata.filter_metadata`¶

Step 19 · `get_preprocessor_with_audio_normalization`¶

Step 20 · `UuidSampler`¶

Step 21 · `KintsugiMetadata.getitem`¶

Step 22 · `KintsugiAudioDataset.getitem`¶

Step 23 · `get_window_starts`¶

Step 24 · `broadcast_collate`¶

Step 25 · `MultiTaskOrdinalRegressionPLModule.training_step`¶

Step 26 · `MultiTaskOrdinalRegressionPLModule.step`¶

Step 27 · `DepAnxClassifier.forward`¶

Step 28 · `DepAnxClassifier.forward_backbone`¶

Step 29 · `WhisperTorchFeatureExtractor.forward`¶

Step 30 · `WhisperBackbone.forward`¶

Step 31 · `MultitaskHead.forward`¶

Step 32 · `OrdinalRegressionPLModule.step`¶

Step 33 · `OrdinalRegressionPLModule.compute_and_log_loss`¶

Step 34 · `average_and_variance_tensor_in_segments`¶

Step 35 · `CORALLoss.forward`¶

Step 36 · `MultiTaskOrdinalRegressionPLModule.on_validation_test_epoch_end`¶

Step 37 · `OrdinalRegressionPLModule.on_validation_test_epoch_end`¶

Step 38 · `OrdinalRegressionPLModule.quantize_labels`¶

Step 39 · `IndetSnSpArray.build`¶

Step 40 · `OptimalOrdinalThresholdingViaDynamicProgramming.tune_thresholds`¶

Step 41 · `OrdinalRegressionPLModule.log_ordinal_metrics`¶

Step 42 · `DistributedModelCheckpoint.on_validation_end`¶

Step 43 · `_log_datasets_to_wandb`¶

Step 44 · `get_dataframe_with_scores`¶