diff --git a/climanet/predict.py b/climanet/predict.py index 84845ef..aa457d9 100644 --- a/climanet/predict.py +++ b/climanet/predict.py @@ -59,6 +59,7 @@ def predict_monthly_var( run_dir: str = ".", verbose: bool = True, dataloader_num_workers: int = 2, + store_logs: bool = True, ): """ Predicts monthly variable values using a trained model and a provided dataset. @@ -105,7 +106,8 @@ def predict_monthly_var( all_predictions = torch.empty(len(dataset), M, H, W, device=device) # Set up logging - writer = setup_logging(run_dir) + if store_logs: + writer = setup_logging(run_dir) with torch.inference_mode(): idx = 0 @@ -140,13 +142,16 @@ def predict_monthly_var( f"Processed batch {i + 1}/{len(dataloader)}, with loss: {loss.item():.4f}" ) - writer.add_scalar("Progress/Batch", i + 1, idx) + if store_logs: + writer.add_scalar("Progress/Batch", i + 1, idx) average_loss = average_loss.item() / len(dataloader) if verbose: print(f"Average loss over all batches: {average_loss:.4f}") - writer.add_scalar("Loss/Average", average_loss) + + if store_logs: + writer.add_scalar("Loss/Average", average_loss) if return_numpy: all_predictions = all_predictions.cpu().numpy() @@ -159,10 +164,12 @@ def predict_monthly_var( if verbose: print(f"Predictions saved to '{run_dir}'") - writer.add_text("Info", f"Predictions saved to '{run_dir}'") + if store_logs: + writer.add_text("Info", f"Predictions saved to '{run_dir}'") # Close the writer when done - writer.close() + if store_logs: + writer.close() if return_loss: all_predictions = (all_predictions, average_loss) diff --git a/climanet/train.py b/climanet/train.py index a6e0055..8e2c4d0 100644 --- a/climanet/train.py +++ b/climanet/train.py @@ -1,3 +1,4 @@ +import tempfile from pathlib import Path import torch @@ -41,6 +42,7 @@ def train_monthly_model( dataloader_num_workers: int = 2, verbose_epoch_interval: int = 20, tune_checkpoint: bool = False, + store_logs: bool = True, ): """Train the model to predict monthly data from daily data. Args: @@ -75,7 +77,8 @@ def train_monthly_model( ) # Set up logging - writer = setup_logging(run_dir) + if store_logs: + writer = setup_logging(run_dir) # Set the optimizer optimizer = torch.optim.AdamW( @@ -133,7 +136,10 @@ def train_monthly_model( # Calculate average epoch loss avg_train_loss = epoch_loss.item() / (i + 1) - writer.add_scalar("Loss/train", avg_train_loss, epoch) + + if store_logs: + writer.add_scalar("Loss/train", avg_train_loss, epoch) + avg_epoch_loss = avg_train_loss # Initially use training loss # Validation loss (optional) @@ -150,10 +156,13 @@ def train_monthly_model( verbose=False, run_dir=run_dir, dataloader_num_workers=dataloader_num_workers, + store_logs=False, ) - writer.add_scalar("Loss/validation", avg_val_loss, epoch) avg_epoch_loss = avg_val_loss # Use validation loss if exists + if store_logs: + writer.add_scalar("Loss/validation", avg_val_loss, epoch) + if verbose and epoch % verbose_epoch_interval == 0: gap = avg_val_loss - avg_train_loss print(f"Epoch {epoch}: gap between train and val loss: {gap:.6f}") @@ -173,7 +182,8 @@ def train_monthly_model( counter += 1 # Log to TensorBoard - writer.add_scalar("Loss/best", best_loss, epoch) + if store_logs: + writer.add_scalar("Loss/best", best_loss, epoch) if verbose and epoch % 20 == 0: print(f"Epoch {epoch}: best_loss = {best_loss:.6f}") @@ -181,7 +191,8 @@ def train_monthly_model( # Only stop if LR is at minimum AND no improvement current_lr = optimizer.param_groups[0]["lr"] if counter >= patience and current_lr <= scheduler.min_lrs[0]: - writer.add_text("Training", f"Early stop at epoch {epoch}", epoch) + if store_logs: + writer.add_text("Training", f"Early stop at epoch {epoch}", epoch) break # Restore best model @@ -189,14 +200,18 @@ def train_monthly_model( model.load_state_dict(best_state_dict) if tune_checkpoint: - # Save the model and optimizer state for Ray Tune checkpointing - save_model(model, optimizer, run_dir, filename="checkpoint.pt", verbose=False) - tune.report( - {"loss": best_loss}, checkpoint=tune.Checkpoint.from_directory(run_dir) - ) + with tempfile.TemporaryDirectory() as checkpoint_dir: + checkpoint_path = Path(checkpoint_dir) + + # Save the model and optimizer state for Ray Tune checkpointing + save_model(model, optimizer, checkpoint_path, filename="checkpoint.pt", verbose=False) + tune.report( + {"loss": best_loss}, checkpoint=tune.Checkpoint.from_directory(checkpoint_path) + ) # Close the writer when done - writer.close() + if store_logs: + writer.close() if verbose: print(f"Training complete. Best loss: {best_loss:.6f}") diff --git a/climanet/tune.py b/climanet/tune.py index 3180f33..5ea8df5 100644 --- a/climanet/tune.py +++ b/climanet/tune.py @@ -2,6 +2,7 @@ import ray import xarray as xr +from ray.air.config import CheckpointConfig from ray.tune.schedulers import ASHAScheduler from climanet.dataset import STDataset @@ -20,8 +21,15 @@ def _train(tune_config, static_args): num_epoch = static_args["num_epoch"] # dont use ray.put() and ray.get() (i.e. object store) when data is large - train_dataset = tune_data_preparation(static_args["data_config_train"]) - validation_dataset = tune_data_preparation(static_args["data_config_validation"]) + if static_args.get("train_dataset") is not None: + train_dataset = ray.get(static_args["train_dataset"]) + else: + train_dataset = tune_data_preparation(static_args["data_config_train"]) + + if static_args.get("validation_dataset") is not None: + validation_dataset = ray.get(static_args["validation_dataset"]) + else: + validation_dataset = tune_data_preparation(static_args["data_config_validation"]) patch_size = tune_config["patch_size"] overlap = tune_config["overlap"] @@ -61,6 +69,7 @@ def _train(tune_config, static_args): store_model=False, verbose=False, tune_checkpoint=True, + store_logs=False, ) @@ -117,6 +126,13 @@ def run_tune(tune_config: dict, static_args: dict): if Path(experiment_path).exists(): tuner = ray.tune.Tuner.restore( experiment_path, + ray.tune.with_resources( + ray.tune.with_parameters(_train, static_args=static_args), + resources={ + "cpu": static_args["cpu_per_trial"], + "gpu": static_args["gpu_per_trial"], + }, + ), resume_errored=True, ) else: @@ -136,7 +152,15 @@ def run_tune(tune_config: dict, static_args: dict): max_concurrent_trials=static_args["max_concurrent_trials"], ), param_space=tune_config, - run_config=ray.tune.RunConfig(storage_path=static_args["run_dir"], name=experiment_name), + run_config=ray.tune.RunConfig( + storage_path=static_args["run_dir"], + name=experiment_name, + checkpoint_config=CheckpointConfig( + num_to_keep=1, + checkpoint_score_attribute="loss", + checkpoint_score_order="min", + ), + ), ) results = tuner.fit() diff --git a/notebooks/example_tuning.ipynb b/notebooks/example_tuning.ipynb index bedceac..d80aedd 100644 --- a/notebooks/example_tuning.ipynb +++ b/notebooks/example_tuning.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 15, + "execution_count": null, "id": "57546057-2042-42eb-a793-23e77d22965e", "metadata": {}, "outputs": [], @@ -168,8 +168,8 @@ "name": "stderr", "output_type": "stream", "text": [ - "2026-07-23 10:11:28,180\tINFO worker.py:2024 -- Started a local Ray instance.\n", - "\u001b[36m(_train pid=40937)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56/_train_2bdd9_00001_1_patch_size=4_2026-07-23_10-11-56/checkpoint_000000)\n" + "2026-08-03 10:59:39,452\tINFO worker.py:2024 -- Started a local Ray instance.\n", + "\u001b[36m(_train pid=127429)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune/_train_ac0e6_00001_1_patch_size=4_2026-08-03_10-59-44/checkpoint_000000)\n" ] } ], @@ -185,7 +185,8 @@ " \"train_dataset\": ray.put(train_dataset),\n", " \"validation_dataset\": ray.put(validation_dataset),\n", " \"num_epoch\": 1,\n", - " \"max_concurrent_trials\": 2\n", + " \"max_concurrent_trials\": 2,\n", + " \"experiment_name\": \"climanet_tune\",\n", "}\n", "\n", "# parameters to tune\n", @@ -225,9 +226,9 @@ "

Tune Status

\n", " \n", "\n", - "\n", - "\n", - "\n", + "\n", + "\n", + "\n", "\n", "
Current time:2026-07-23 10:12:08
Running for: 00:00:11.76
Memory: 8.6/15.3 GiB
Current time:2026-08-03 10:59:57
Running for: 00:00:13.08
Memory: 9.8/15.3 GiB
\n", " \n", @@ -243,11 +244,11 @@ "

Trial Status

\n", " \n", "\n", - "\n", + "\n", "\n", "\n", - "\n", - "\n", + "\n", + "\n", "\n", "
Trial name status loc patch_size iter total time (s) loss
Trial name status loc patch_size iter total time (s) loss
_train_2bdd9_00000TERMINATED192.168.2.13:40938 2 1 6.458450.384662
_train_2bdd9_00001TERMINATED192.168.2.13:40937 4 1 3.896690.365478
_train_ac0e6_00000TERMINATED192.168.2.13:127430 2 1 6.447040.384662
_train_ac0e6_00001TERMINATED192.168.2.13:127429 4 1 3.764730.365478
\n", " \n", @@ -294,9 +295,15 @@ "name": "stderr", "output_type": "stream", "text": [ - "2026-07-23 10:12:08,344\tINFO tune.py:1007 -- Wrote the latest version of all result files and experiment state to '/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56' in 0.0033s.\n", - "2026-07-23 10:12:08,348\tINFO tune.py:1039 -- Total run time: 11.79 seconds (11.75 seconds for the tuning loop).\n", - "\u001b[36m(_train pid=40938)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56/_train_2bdd9_00000_0_patch_size=2_2026-07-23_10-11-56/checkpoint_000000)\n" + "2026-08-03 10:59:55,339\tWARNING experiment_state.py:209 -- Experiment state snapshotting has been triggered multiple times in the last 5.0 seconds and may become a bottleneck. A snapshot is forced if `CheckpointConfig(num_to_keep)` is set, and a trial has checkpointed >= `num_to_keep` times since the last snapshot.\n", + "You may want to consider increasing the `CheckpointConfig(num_to_keep)` or decreasing the frequency of saving checkpoints.\n", + "You can suppress this warning by setting the environment variable TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S to a smaller value than the current threshold (5.0). Set it to 0 to completely suppress this warning.\n", + "2026-08-03 10:59:57,963\tWARNING experiment_state.py:209 -- Experiment state snapshotting has been triggered multiple times in the last 5.0 seconds and may become a bottleneck. A snapshot is forced if `CheckpointConfig(num_to_keep)` is set, and a trial has checkpointed >= `num_to_keep` times since the last snapshot.\n", + "You may want to consider increasing the `CheckpointConfig(num_to_keep)` or decreasing the frequency of saving checkpoints.\n", + "You can suppress this warning by setting the environment variable TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S to a smaller value than the current threshold (5.0). Set it to 0 to completely suppress this warning.\n", + "2026-08-03 10:59:57,971\tINFO tune.py:1007 -- Wrote the latest version of all result files and experiment state to '/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune' in 0.0066s.\n", + "2026-08-03 10:59:57,986\tINFO tune.py:1039 -- Total run time: 13.11 seconds (13.07 seconds for the tuning loop).\n", + "\u001b[36m(_train pid=127430)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune/_train_ac0e6_00000_0_patch_size=2_2026-08-03_10-59-44/checkpoint_000000)\n" ] } ], @@ -325,9 +332,9 @@ "text/plain": [ "Result(\n", " metrics={'loss': 0.36547836661338806},\n", - " path='/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56/_train_2bdd9_00001_1_patch_size=4_2026-07-23_10-11-56',\n", + " path='/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune/_train_ac0e6_00001_1_patch_size=4_2026-08-03_10-59-44',\n", " filesystem='local',\n", - " checkpoint=Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56/_train_2bdd9_00001_1_patch_size=4_2026-07-23_10-11-56/checkpoint_000000)\n", + " checkpoint=Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune/_train_ac0e6_00001_1_patch_size=4_2026-08-03_10-59-44/checkpoint_000000)\n", ")" ] }, @@ -350,10 +357,18 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": 10, "id": "66645a30-89d7-459b-9fcc-e03fa4d34b71", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2026-08-03 11:01:08,865\tINFO worker.py:2024 -- Started a local Ray instance.\n" + ] + } + ], "source": [ "# find the path to best model\n", "if not ray.is_initialized():\n", @@ -368,7 +383,7 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": 11, "id": "67274359-2e8a-485f-a56a-a3831bd5f9cb", "metadata": {}, "outputs": [], @@ -394,7 +409,7 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 12, "id": "65c8bab8-dc84-46ee-a249-24b7dc24f991", "metadata": {}, "outputs": [ @@ -404,7 +419,7 @@ "0.3733350435892741" ] }, - "execution_count": 20, + "execution_count": 12, "metadata": {}, "output_type": "execute_result" }