Add SLURM training scripts and fix Comet.ml logging

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",

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"
     )

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,9 @@ 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:
     """Log validation metrics to Comet.ml.
 
@@ -364,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
@@ -376,9 +382,33 @@ 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:
+        metrics["val/gp_nll"] = val_gp_nll
     _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,

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"

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 <config_file> [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

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

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,27 +46,28 @@ 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
 lr: 5e-4
 final_lr: 1e-5
 weight_decay: 0.0001
 gradient_accumulation_steps: 1
-epochs: 10
-frac_warmup_steps: 0.1
+epochs: 200
+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', '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', 'run2']
+comet_project: multiplex-image-model
+comet_workspace: micha-zmys-owski
+comet_api_key: null  # set via COMET_API_KEY env var

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,20 +261,24 @@ 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,
         "val_mae": val_mae,
         "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()