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/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/multiplex_model/utils/train_logging.py b/multiplex_model/utils/train_logging.py index 83e7ddd..48bc4fd 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,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, 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..7da8d30 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 @@ -53,19 +54,20 @@ 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 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 d601671..5f503e1 100644 --- a/train_masked_model_gp.py +++ b/train_masked_model_gp.py @@ -49,9 +49,11 @@ get_scheduler_with_warmup, init_experiment, log_training_metrics, + log_validation_batch_metrics, log_validation_images, log_validation_metrics, plot_reconstructs_with_masks, + plot_reconstructs_with_uncertainty, ) @@ -242,6 +244,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() @@ -294,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( @@ -326,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: @@ -367,6 +381,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) @@ -376,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, @@ -389,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, } @@ -406,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() @@ -552,10 +592,26 @@ 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 - total_steps = ( - 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 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: + 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 @@ -578,6 +634,10 @@ def test_masked_gp( type="cosine", ) + 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"]) + # Initialize experiment tracking comet_config = config.model_dump() comet_config.update({ @@ -592,18 +652,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,