From 9b343c8776df3bdc1e99f45a9e14ad80c9ccd287 Mon Sep 17 00:00:00 2001 From: mzmyslowski Date: Tue, 17 Mar 2026 11:22:45 +0100 Subject: [PATCH 1/5] Add SLURM training scripts and fix Comet.ml logging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add train.sh SLURM batch script for standard and GP training - Add setup_venv.sh for uv-based venv setup on remote server - Fix Comet.ml empty tags error (skip add_tags when tags=[]) - Add val_standard_nll and val_gp_nll to log_validation_metrics - Fix train_masked_config.yaml and train_masked_gp_config.yaml Comet placeholders - Set gp_lengthscale 0.1→5.0, lambda_gp 0.0→0.1, batch_size 1→8, use_kronecker_gp=true in GP config Co-Authored-By: Claude Sonnet 4.6 --- multiplex_model/utils/train_logging.py | 10 ++++++- setup_venv.sh | 19 +++++++++++++ train.sh | 38 ++++++++++++++++++++++++++ train_masked_config.yaml | 6 ++-- train_masked_gp_config.yaml | 15 +++++----- 5 files changed, 77 insertions(+), 11 deletions(-) create mode 100755 setup_venv.sh create mode 100755 train.sh diff --git a/multiplex_model/utils/train_logging.py b/multiplex_model/utils/train_logging.py index 83e7ddd..b219030 100644 --- a/multiplex_model/utils/train_logging.py +++ b/multiplex_model/utils/train_logging.py @@ -301,7 +301,9 @@ def init_experiment(config: dict[str, Any]) -> None: print(f"Run name: {run_name}") _experiment.set_name(run_name) - _experiment.add_tags(config.get("tags", [])) + tags = config.get("tags", []) + if tags: + _experiment.add_tags(tags) _experiment.log_parameters(config) @@ -354,6 +356,8 @@ def log_validation_metrics( latent_rankme: float, epoch: int, variance_mae_correlation: float | None = None, + val_standard_nll: float | None = None, + val_gp_nll: float | None = None, ) -> None: """Log validation metrics to Comet.ml. @@ -376,6 +380,10 @@ def log_validation_metrics( } if variance_mae_correlation is not None: metrics["val/variance_mae_correlation"] = variance_mae_correlation + if val_standard_nll is not None: + metrics["val/standard_nll"] = val_standard_nll + if val_gp_nll is not None: + metrics["val/gp_nll"] = val_gp_nll _experiment.log_metrics(metrics, epoch=epoch) diff --git a/setup_venv.sh b/setup_venv.sh new file mode 100755 index 0000000..94a0ed5 --- /dev/null +++ b/setup_venv.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Run this once on bury/szary to set up the virtual environment +set -e + +# Install uv if not present +if ! command -v uv &> /dev/null; then + curl -LsSf https://astral.sh/uv/install.sh | sh + source "$HOME/.local/bin/env" +fi + +export PATH="$HOME/.local/bin:$PATH" + +cd "$(dirname "$0")" + +uv venv ~/venv +source ~/venv/bin/activate +uv pip install -e ".[dev]" + +echo "Venv ready at ~/venv" diff --git a/train.sh b/train.sh new file mode 100755 index 0000000..0f56732 --- /dev/null +++ b/train.sh @@ -0,0 +1,38 @@ +#!/bin/bash +#SBATCH --partition=common +#SBATCH --qos=mzmyslowski +#SBATCH --nodelist=szary +#SBATCH --cpus-per-task=8 +#SBATCH --mem=50G +#SBATCH --gres=gpu:1 +#SBATCH --time=24:00:00 +#SBATCH --job-name=train +#SBATCH --output=logs/train_%j.out +#SBATCH --error=logs/train_%j.err + +set -e + +if [ -z "$1" ]; then + echo "Usage: sbatch train.sh [gp]" + echo " config_file: path to YAML config" + echo " gp: pass 'gp' as second arg to use GP training script" + exit 1 +fi + +config_file=$1 +use_gp=${2:-""} + +mkdir -p logs + +# Set COMET_API_KEY in your environment or ~/.bashrc before submitting +# export COMET_API_KEY=your_key_here + +source ~/venv/bin/activate + +if [ "$use_gp" = "gp" ]; then + echo "Starting GP training with config: $config_file" + python3 train_masked_model_gp.py "$config_file" +else + echo "Starting standard training with config: $config_file" + python3 train_masked_model.py "$config_file" +fi diff --git a/train_masked_config.yaml b/train_masked_config.yaml index ddb2db7..fa275ca 100644 --- a/train_masked_config.yaml +++ b/train_masked_config.yaml @@ -47,7 +47,7 @@ save_checkpoint_freq: 5 beta: 1.0 # Comet.ml logging configuration -tags: [...] -comet_project: ... -comet_workspace: null # optional, can also be set via COMET_WORKSPACE env var +tags: [] +comet_project: multiplex-image-model +comet_workspace: micha-zmys-owski # optional, can also be set via COMET_WORKSPACE env var comet_api_key: null # optional, can also be set via COMET_API_KEY env var diff --git a/train_masked_gp_config.yaml b/train_masked_gp_config.yaml index 206f063..2c97c53 100644 --- a/train_masked_gp_config.yaml +++ b/train_masked_gp_config.yaml @@ -5,9 +5,10 @@ # GP LOSS CONFIGURATION # ============================================================================ use_gp_loss: true # Enable/disable GP loss -lambda_gp: 0.0 # Weight for GP loss (0.0 = only standard, 1.0 = only GP) +use_kronecker_gp: true # Use Kronecker (~40x faster than CG) +lambda_gp: 0.1 # Weight for GP loss (0.0 = only standard, 1.0 = only GP) gp_kernel_jitter: 1e-2 # Diagonal noise for numerical stability -gp_lengthscale: 0.1 # Spatial correlation length scale +gp_lengthscale: 5.0 # Spatial correlation length scale gp_max_cg_iterations: 50 # Max conjugate gradient iterations gp_downscale_factor: 1 # Spatial downsampling (1=none, 2=half, 4=quarter) gp_learn_lengthscale: true # Whether to learn kernel lengthscale @@ -45,7 +46,7 @@ panel_config: configs/all_panels_config.yaml tokenizer_config: configs/all_markers_tokenizer.yaml input_image_size: [112, 112] num_workers: 8 -batch_size: 1 +batch_size: 8 # Training configuration device: cuda @@ -65,7 +66,7 @@ save_checkpoint_freq: 5 beta: 0.5 # Comet.ml logging configuration -tags: ['SZARY', 'GP', 'times', 'lambda 0.1'] -comet_project: ... -comet_workspace: ... # optional, can also be set via COMET_WORKSPACE env var -comet_api_key: ... # optional, can also be set via COMET_API_KEY env var +tags: ['SZARY', 'GP', 'kronecker', 'lambda 0.1'] +comet_project: multiplex-image-model +comet_workspace: micha-zmys-owski +comet_api_key: null # set via COMET_API_KEY env var From 9b1baea2651b5fed4a11b300e7873c1aa5c00f2d Mon Sep 17 00:00:00 2001 From: mzmyslowski Date: Tue, 17 Mar 2026 13:24:04 +0100 Subject: [PATCH 2/5] Fix scheduler resumption from checkpoint and set 100 epochs for GP training MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Save total_steps in checkpoint so scheduler can be reconstructed with identical warmup/annealing boundaries when resuming. Previously resuming with different epochs config would miscalculate LR schedule. Also bump epochs 10→100 in GP config. Co-Authored-By: Claude Sonnet 4.6 --- train_masked_gp_config.yaml | 2 +- train_masked_model_gp.py | 33 ++++++++++++++++++++------------- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/train_masked_gp_config.yaml b/train_masked_gp_config.yaml index 2c97c53..4351c15 100644 --- a/train_masked_gp_config.yaml +++ b/train_masked_gp_config.yaml @@ -54,7 +54,7 @@ lr: 5e-4 final_lr: 1e-5 weight_decay: 0.0001 gradient_accumulation_steps: 1 -epochs: 10 +epochs: 100 frac_warmup_steps: 0.1 min_channels_frac: 0.75 spatial_masking_ratio: 0.6 diff --git a/train_masked_model_gp.py b/train_masked_model_gp.py index d601671..863c172 100644 --- a/train_masked_model_gp.py +++ b/train_masked_model_gp.py @@ -242,6 +242,7 @@ def train_masked_gp( "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "epoch": epoch, + "total_steps": total_steps, } if gp_covariance_module is not None: checkpoint["gp_covariance_state_dict"] = gp_covariance_module.state_dict() @@ -552,9 +553,23 @@ def test_masked_gp( device=device, ) + # Load checkpoint early to recover total_steps for scheduler reconstruction + start_epoch = 0 + checkpoint = None + if config.resolve_checkpoint(): + print(f"Loading model from checkpoint: {config.from_checkpoint}") + checkpoint = torch.load(config.from_checkpoint, map_location=device) + model.load_state_dict(checkpoint["model_state_dict"]) + if gp_covariance_module is not None and "gp_covariance_state_dict" in checkpoint: + gp_covariance_module.load_state_dict(checkpoint["gp_covariance_state_dict"]) + start_epoch = checkpoint["epoch"] + 1 + # Optimizer and scheduler + # When resuming, use saved total_steps so scheduler boundaries match original run total_steps = ( - len(train_dataloader) * config.epochs // config.gradient_accumulation_steps + checkpoint["total_steps"] + if checkpoint is not None and "total_steps" in checkpoint + else len(train_dataloader) * config.epochs // config.gradient_accumulation_steps ) num_warmup_steps = int(total_steps * config.frac_warmup_steps) num_annealing_steps = total_steps - num_warmup_steps @@ -578,6 +593,10 @@ def test_masked_gp( type="cosine", ) + if checkpoint is not None: + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) + # Initialize experiment tracking comet_config = config.model_dump() comet_config.update({ @@ -592,18 +611,6 @@ def test_masked_gp( }) init_experiment(comet_config) - # Load checkpoint if specified - start_epoch = 0 - if config.resolve_checkpoint(): - print(f"Loading model from checkpoint: {config.from_checkpoint}") - checkpoint = torch.load(config.from_checkpoint, map_location=device) - model.load_state_dict(checkpoint["model_state_dict"]) - optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) - scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) - if gp_covariance_module is not None and "gp_covariance_state_dict" in checkpoint: - gp_covariance_module.load_state_dict(checkpoint["gp_covariance_state_dict"]) - start_epoch = checkpoint["epoch"] + 1 - # Train the model train_masked_gp( model, From 7a58ad4a2d411e70cc1b7b83050162e4bcbfa08a Mon Sep 17 00:00:00 2001 From: mzmyslowski Date: Mon, 23 Mar 2026 21:26:57 +0100 Subject: [PATCH 3/5] Add reset_lr_schedule flag for fresh cosine cycle on resumption When training for another N epochs from a fully-converged checkpoint, the saved scheduler state has LR near zero. reset_lr_schedule=true ignores checkpoint optimizer/scheduler state and starts a fresh cosine cycle from the trained weights. Co-Authored-By: Claude Sonnet 4.6 --- multiplex_model/utils/configuration.py | 4 +++ train_masked_gp_config.yaml | 7 +++--- train_masked_model_gp.py | 35 ++++++++++++++++++++------ 3 files changed, 36 insertions(+), 10 deletions(-) diff --git a/multiplex_model/utils/configuration.py b/multiplex_model/utils/configuration.py index 736ff46..55724fb 100644 --- a/multiplex_model/utils/configuration.py +++ b/multiplex_model/utils/configuration.py @@ -256,6 +256,10 @@ class TrainingConfig(BaseModel): None, description="Path to checkpoint to resume from. Use 'last' to load last checkpoint if available", ) + reset_lr_schedule: bool = Field( + False, + description="When resuming, ignore checkpoint's scheduler/optimizer state and start a fresh LR schedule", + ) checkpoints_dir: str = Field( "checkpoints", description="Directory to save checkpoints" ) diff --git a/train_masked_gp_config.yaml b/train_masked_gp_config.yaml index 4351c15..747d55c 100644 --- a/train_masked_gp_config.yaml +++ b/train_masked_gp_config.yaml @@ -55,18 +55,19 @@ final_lr: 1e-5 weight_decay: 0.0001 gradient_accumulation_steps: 1 epochs: 100 -frac_warmup_steps: 0.1 +frac_warmup_steps: 0.01 min_channels_frac: 0.75 spatial_masking_ratio: 0.6 fully_masked_channels_max_frac: 0.5 mask_patch_size: 8 -from_checkpoint: null +from_checkpoint: checkpoints/final_model-ImVs-12.pth +reset_lr_schedule: true # fresh cosine cycle from trained weights checkpoints_dir: checkpoints save_checkpoint_freq: 5 beta: 0.5 # Comet.ml logging configuration -tags: ['SZARY', 'GP', 'kronecker', 'lambda 0.1'] +tags: ['SZARY', 'GP', 'kronecker', 'lambda 0.1', 'run2'] comet_project: multiplex-image-model comet_workspace: micha-zmys-owski comet_api_key: null # set via COMET_API_KEY env var diff --git a/train_masked_model_gp.py b/train_masked_model_gp.py index 863c172..b56ffb4 100644 --- a/train_masked_model_gp.py +++ b/train_masked_model_gp.py @@ -52,6 +52,7 @@ log_validation_images, log_validation_metrics, plot_reconstructs_with_masks, + plot_reconstructs_with_uncertainty, ) @@ -368,6 +369,26 @@ def test_masked_gp( masked_channels_names=masked_channels_names, img_idx=idx, ) + + sigma = torch.exp(0.5 * logvar) + uncertainty_img = plot_reconstructs_with_uncertainty( + img, + mi, + sigma, + channel_ids, + unactive_channels, + markers_names_map=marker_names_map, + ncols=9, + ) + log_validation_images( + fig=uncertainty_img, + panel_idx=panel_idx[0], + img_path=img_path[0], + epoch=epoch, + masked_channels_names=masked_channels_names, + img_idx=idx, + name_suffix="_sigma", + ) plt.close("all") val_loss = running_loss / len(test_dataloader) @@ -565,12 +586,12 @@ def test_masked_gp( start_epoch = checkpoint["epoch"] + 1 # Optimizer and scheduler - # When resuming, use saved total_steps so scheduler boundaries match original run - total_steps = ( - checkpoint["total_steps"] - if checkpoint is not None and "total_steps" in checkpoint - else len(train_dataloader) * config.epochs // config.gradient_accumulation_steps - ) + # When resuming normally, use saved total_steps so scheduler boundaries match original run. + # When reset_lr_schedule=True, recalculate from config.epochs for a fresh cosine cycle. + if checkpoint is not None and "total_steps" in checkpoint and not config.reset_lr_schedule: + total_steps = checkpoint["total_steps"] + else: + total_steps = len(train_dataloader) * config.epochs // config.gradient_accumulation_steps num_warmup_steps = int(total_steps * config.frac_warmup_steps) num_annealing_steps = total_steps - num_warmup_steps @@ -593,7 +614,7 @@ def test_masked_gp( type="cosine", ) - if checkpoint is not None: + if checkpoint is not None and not config.reset_lr_schedule: optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) From a5a42a6e1bedb16d7504293365ff814c27380d15 Mon Sep 17 00:00:00 2001 From: mzmyslowski Date: Mon, 23 Mar 2026 21:30:44 +0100 Subject: [PATCH 4/5] Fix empty training loop and total_steps mismatch on fresh LR reset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two bugs when using reset_lr_schedule=True: 1. epochs:100 + start_epoch=100 → range(100,100) empty, no training Fix: epochs:200 so range(100,200) = 100 new epochs 2. total_steps calculated from config.epochs (200) but only 100 epochs will run → cosine schedule only half-completed at end of run Fix: use remaining_epochs = config.epochs - start_epoch for total_steps Co-Authored-By: Claude Sonnet 4.6 --- train_masked_gp_config.yaml | 2 +- train_masked_model_gp.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/train_masked_gp_config.yaml b/train_masked_gp_config.yaml index 747d55c..7da8d30 100644 --- a/train_masked_gp_config.yaml +++ b/train_masked_gp_config.yaml @@ -54,7 +54,7 @@ lr: 5e-4 final_lr: 1e-5 weight_decay: 0.0001 gradient_accumulation_steps: 1 -epochs: 100 +epochs: 200 frac_warmup_steps: 0.01 min_channels_frac: 0.75 spatial_masking_ratio: 0.6 diff --git a/train_masked_model_gp.py b/train_masked_model_gp.py index b56ffb4..2be4c9d 100644 --- a/train_masked_model_gp.py +++ b/train_masked_model_gp.py @@ -587,11 +587,13 @@ def test_masked_gp( # Optimizer and scheduler # When resuming normally, use saved total_steps so scheduler boundaries match original run. - # When reset_lr_schedule=True, recalculate from config.epochs for a fresh cosine cycle. + # When reset_lr_schedule=True, recalculate from remaining epochs for a fresh cosine cycle + # that covers exactly the new training run (config.epochs - start_epoch epochs). if checkpoint is not None and "total_steps" in checkpoint and not config.reset_lr_schedule: total_steps = checkpoint["total_steps"] else: - total_steps = len(train_dataloader) * config.epochs // config.gradient_accumulation_steps + remaining_epochs = config.epochs - start_epoch + total_steps = len(train_dataloader) * remaining_epochs // config.gradient_accumulation_steps num_warmup_steps = int(total_steps * config.frac_warmup_steps) num_annealing_steps = total_steps - num_warmup_steps From 0fc93856b1dd4e72c3b2994e4db5a34d8f541711 Mon Sep 17 00:00:00 2001 From: mzmyslowski Date: Fri, 3 Apr 2026 07:55:13 +0200 Subject: [PATCH 5/5] feat(logging): add variance-MSE correlation metrics to validation logging --- multiplex_model/utils/__init__.py | 2 ++ multiplex_model/utils/train_logging.py | 22 ++++++++++++++++++++++ train_masked_model.py | 21 +++++++++++++++++++-- train_masked_model_gp.py | 20 +++++++++++++++++++- 4 files changed, 62 insertions(+), 3 deletions(-) diff --git a/multiplex_model/utils/__init__.py b/multiplex_model/utils/__init__.py index 9aeb72d..6200bd2 100644 --- a/multiplex_model/utils/__init__.py +++ b/multiplex_model/utils/__init__.py @@ -25,6 +25,7 @@ get_run_name, init_experiment, log_training_metrics, + log_validation_batch_metrics, log_validation_images, log_validation_metrics, plot_reconstructs_with_masks, @@ -46,6 +47,7 @@ "plot_reconstructs_with_masks", "init_experiment", "log_training_metrics", + "log_validation_batch_metrics", "log_validation_metrics", "log_validation_images", "get_run_name", diff --git a/multiplex_model/utils/train_logging.py b/multiplex_model/utils/train_logging.py index b219030..48bc4fd 100644 --- a/multiplex_model/utils/train_logging.py +++ b/multiplex_model/utils/train_logging.py @@ -356,6 +356,7 @@ def log_validation_metrics( latent_rankme: float, epoch: int, variance_mae_correlation: float | None = None, + variance_mse_correlation: float | None = None, val_standard_nll: float | None = None, val_gp_nll: float | None = None, ) -> None: @@ -368,6 +369,7 @@ def log_validation_metrics( latent_rankme (float): RankMe metric for latent representations epoch (int): Current epoch number variance_mae_correlation (Optional[float]): Pearson correlation between predicted variances and MAEs per channel + variance_mse_correlation (Optional[float]): Pearson correlation between predicted variances and MSEs per channel """ if _experiment is None: return @@ -380,6 +382,8 @@ def log_validation_metrics( } if variance_mae_correlation is not None: metrics["val/variance_mae_correlation"] = variance_mae_correlation + if variance_mse_correlation is not None: + metrics["val/variance_mse_correlation"] = variance_mse_correlation if val_standard_nll is not None: metrics["val/standard_nll"] = val_standard_nll if val_gp_nll is not None: @@ -387,6 +391,24 @@ def log_validation_metrics( _experiment.log_metrics(metrics, epoch=epoch) +def log_validation_batch_metrics( + variance_mse_correlation_per_batch: float, + step: int, +) -> None: + """Log per-batch validation metrics to Comet.ml. + + Args: + variance_mse_correlation_per_batch (float): Pearson correlation between predicted variances and MSEs per channel for a single batch + step (int): Global step number + """ + if _experiment is None: + return + _experiment.log_metrics( + {"val/variance_mse_correlation_per_batch": variance_mse_correlation_per_batch}, + step=step, + ) + + def log_validation_images( fig: plt.Figure, panel_idx: int, diff --git a/train_masked_model.py b/train_masked_model.py index 0d230ee..32c61f2 100644 --- a/train_masked_model.py +++ b/train_masked_model.py @@ -32,6 +32,7 @@ get_scheduler_with_warmup, init_experiment, log_training_metrics, + log_validation_batch_metrics, log_validation_images, log_validation_metrics, plot_reconstructs_with_masks, @@ -173,6 +174,7 @@ def test_masked( all_latents = [] all_channel_variances = [] all_channel_maes = [] + all_channel_mses = [] with torch.no_grad(): for idx, (img, channel_ids, panel_idx, img_path) in enumerate( @@ -207,8 +209,18 @@ def test_masked( dim=(0, 2, 3) ) # Mean variance per channel mae_per_channel = torch.abs(img - mi).mean(dim=(0, 2, 3)) # MAE per channel + mse_per_channel = torch.square(img - mi).mean(dim=(0, 2, 3)) # MSE per channel all_channel_variances.append(variance_per_channel.cpu()) all_channel_maes.append(mae_per_channel.cpu()) + all_channel_mses.append(mse_per_channel.cpu()) + + batch_var_mse_corr = torch.corrcoef( + torch.stack([variance_per_channel.cpu(), mse_per_channel.cpu()]) + )[0, 1].item() + log_validation_batch_metrics( + variance_mse_correlation_per_batch=batch_var_mse_corr, + step=epoch * len(test_dataloader) + idx, + ) loss = nll_loss(img, mi, logvar) running_loss += loss.item() @@ -249,13 +261,16 @@ def test_masked( all_latents = torch.cat(all_latents) rankme = RankMe(all_latents) - # Calculate Pearson correlation between predicted variances and MAEs per channel + # Calculate Pearson correlation between predicted variances and MAEs/MSEs per channel all_channel_variances = torch.cat(all_channel_variances) all_channel_maes = torch.cat(all_channel_maes) - # Calculate Pearson correlation using flattened data across all batches + all_channel_mses = torch.cat(all_channel_mses) variance_mae_corr = torch.corrcoef( torch.stack([all_channel_variances.flatten(), all_channel_maes.flatten()]) )[0, 1].item() + variance_mse_corr = torch.corrcoef( + torch.stack([all_channel_variances.flatten(), all_channel_mses.flatten()]) + )[0, 1].item() val_metrics = { "val_loss": val_loss, @@ -263,6 +278,7 @@ def test_masked( "val_mse": val_mse, "latent_rankme": rankme, "variance_mae_correlation": variance_mae_corr, + "variance_mse_correlation": variance_mse_corr, "epoch": epoch, } @@ -273,6 +289,7 @@ def test_masked( print(f"MAE: {val_mae:.6f}") print(f"MSE: {val_mse:.6f}") print(f"Pearson MAE vs Var: {variance_mae_corr:.4f}") + print(f"Pearson MSE vs Var: {variance_mse_corr:.4f}") print("=" * 90) print() diff --git a/train_masked_model_gp.py b/train_masked_model_gp.py index 2be4c9d..5f503e1 100644 --- a/train_masked_model_gp.py +++ b/train_masked_model_gp.py @@ -49,6 +49,7 @@ get_scheduler_with_warmup, init_experiment, log_training_metrics, + log_validation_batch_metrics, log_validation_images, log_validation_metrics, plot_reconstructs_with_masks, @@ -296,6 +297,7 @@ def test_masked_gp( all_latents = [] all_channel_variances = [] all_channel_maes = [] + all_channel_mses = [] with torch.no_grad(): for idx, (img, channel_ids, panel_idx, img_path) in enumerate( @@ -328,8 +330,18 @@ def test_masked_gp( # Per-channel statistics variance_per_channel = torch.exp(logvar).mean(dim=(0, 2, 3)) mae_per_channel = torch.abs(img - mi).mean(dim=(0, 2, 3)) + mse_per_channel = torch.square(img - mi).mean(dim=(0, 2, 3)) all_channel_variances.append(variance_per_channel.cpu()) all_channel_maes.append(mae_per_channel.cpu()) + all_channel_mses.append(mse_per_channel.cpu()) + + batch_var_mse_corr = torch.corrcoef( + torch.stack([variance_per_channel.cpu(), mse_per_channel.cpu()]) + )[0, 1].item() + log_validation_batch_metrics( + variance_mse_correlation_per_batch=batch_var_mse_corr, + step=epoch * len(test_dataloader) + idx, + ) # Compute loss if use_gp_loss and gp_loss_fn is not None: @@ -398,12 +410,16 @@ def test_masked_gp( all_latents = torch.cat(all_latents) rankme = RankMe(all_latents) - # Variance-MAE correlation + # Variance-MAE/MSE correlation all_channel_variances = torch.cat(all_channel_variances) all_channel_maes = torch.cat(all_channel_maes) + all_channel_mses = torch.cat(all_channel_mses) variance_mae_corr = torch.corrcoef( torch.stack([all_channel_variances.flatten(), all_channel_maes.flatten()]) )[0, 1].item() + variance_mse_corr = torch.corrcoef( + torch.stack([all_channel_variances.flatten(), all_channel_mses.flatten()]) + )[0, 1].item() val_metrics = { "val_loss": val_loss, @@ -411,6 +427,7 @@ def test_masked_gp( "val_mse": val_mse, "latent_rankme": rankme, "variance_mae_correlation": variance_mae_corr, + "variance_mse_correlation": variance_mse_corr, "epoch": epoch, } @@ -428,6 +445,7 @@ def test_masked_gp( print(f"MAE: {val_mae:.6f}") print(f"MSE: {val_mse:.6f}") print(f"Pearson MAE vs Var: {variance_mae_corr:.4f}") + print(f"Pearson MSE vs Var: {variance_mse_corr:.4f}") print("=" * 90) print()