wehub-resource-sync
222 lines
8.0 KiB
Python
222 lines
8.0 KiB
Python
# Example demonstrating how to use SHOULD_CHECKPOINT in a tuner callback
|
|
# for smart checkpointing logic. This shows how to trigger checkpointing from
|
|
# callbacks based on training progress rather than fixed intervals.
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import time
|
|
|
|
from ray import tune
|
|
from ray.tune import Callback
|
|
from ray.tune.result import SHOULD_CHECKPOINT
|
|
|
|
# Hint: SHOULD_CHECKPOINT is an alias of the string "should_checkpoint"
|
|
|
|
|
|
# Some dummy function
|
|
def evaluation_fn(step, width, height):
|
|
time.sleep(0.1)
|
|
return (0.1 + width * step / 100) ** (-1) + height * 0.1
|
|
|
|
|
|
class SmartCheckpointCallback(Callback):
|
|
"""Custom callback that triggers checkpointing by updating the result dict.
|
|
|
|
This callback demonstrates checkpointing logic beyond
|
|
simple periodic checkpointing. It checkpoints based on performance improvements
|
|
or when the loss becomes unstable.
|
|
|
|
Args:
|
|
checkpoint_on_improvement: Checkpoint when loss improves significantly
|
|
checkpoint_on_instability: Checkpoint when loss becomes unstable
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
checkpoint_on_improvement: bool = True,
|
|
checkpoint_on_instability: bool = True,
|
|
):
|
|
self.checkpoint_on_improvement = checkpoint_on_improvement
|
|
self.checkpoint_on_instability = checkpoint_on_instability
|
|
self.best_loss_per_trial = {}
|
|
self.recent_losses_per_trial = {}
|
|
|
|
def on_trial_result(self, iteration, trials, trial, result, **info):
|
|
"""Called after receiving a result from the trainable.
|
|
|
|
This hook implements intelligent checkpointing logic:
|
|
1. Checkpoint when we see significant improvement
|
|
2. Checkpoint when loss becomes unstable (variance increases)
|
|
3. Always checkpoint at specific milestones (every 10 steps)
|
|
"""
|
|
trial_id = trial.trial_id
|
|
current_loss = result.get("mean_loss", float("inf"))
|
|
current_step = result.get("iterations", 0)
|
|
|
|
# Initialize tracking for this trial
|
|
if trial_id not in self.best_loss_per_trial:
|
|
self.best_loss_per_trial[trial_id] = float("inf")
|
|
self.recent_losses_per_trial[trial_id] = []
|
|
|
|
should_checkpoint = False
|
|
reason = ""
|
|
|
|
# 1. Checkpoint every 10 steps as a baseline
|
|
if current_step > 0 and current_step % 10 == 0:
|
|
should_checkpoint = True
|
|
reason = f"milestone at step {current_step}"
|
|
|
|
# 2. Checkpoint on significant improvement
|
|
if self.checkpoint_on_improvement:
|
|
if (
|
|
current_loss < self.best_loss_per_trial[trial_id] * 0.9
|
|
): # 10% improvement
|
|
should_checkpoint = True
|
|
reason = f"significant improvement: {current_loss:.4f} < {self.best_loss_per_trial[trial_id]:.4f}"
|
|
self.best_loss_per_trial[trial_id] = current_loss
|
|
|
|
# 3. Checkpoint on instability (high variance in recent losses)
|
|
if self.checkpoint_on_instability and current_step > 5:
|
|
recent_losses = self.recent_losses_per_trial[trial_id]
|
|
recent_losses.append(current_loss)
|
|
if len(recent_losses) > 5:
|
|
recent_losses.pop(0) # Keep only last 5 losses
|
|
|
|
if len(recent_losses) == 5:
|
|
variance = (
|
|
sum((x - sum(recent_losses) / 5) ** 2 for x in recent_losses) / 5
|
|
)
|
|
if variance > 0.1: # High variance threshold
|
|
should_checkpoint = True
|
|
reason = f"instability detected: variance={variance:.4f}"
|
|
else:
|
|
# Track recent losses
|
|
recent_losses = self.recent_losses_per_trial[trial_id]
|
|
recent_losses.append(current_loss)
|
|
if len(recent_losses) > 5:
|
|
recent_losses.pop(0)
|
|
|
|
if should_checkpoint:
|
|
print(
|
|
f"Callback requesting checkpoint for trial {trial_id} at step {current_step}: {reason}"
|
|
)
|
|
result[SHOULD_CHECKPOINT] = True
|
|
|
|
|
|
class OptimizationTrainable(tune.Trainable):
|
|
"""A simple trainable that demonstrates automatic checkpointing with callbacks"""
|
|
|
|
def setup(self, config):
|
|
"""Initialize the trainable"""
|
|
self.current_step = 0
|
|
self.width = config["width"]
|
|
self.height = config["height"]
|
|
|
|
def step(self):
|
|
"""Perform one step of training"""
|
|
intermediate_score = evaluation_fn(self.current_step, self.width, self.height)
|
|
self.current_step += 1
|
|
|
|
return {
|
|
"iterations": self.current_step,
|
|
"mean_loss": intermediate_score,
|
|
"step": self.current_step, # For tracking
|
|
}
|
|
|
|
def save_checkpoint(self, checkpoint_dir):
|
|
"""Save checkpoint
|
|
|
|
Called automatically by Tune when SHOULD_CHECKPOINT is in the result
|
|
"""
|
|
checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.json")
|
|
with open(checkpoint_path, "w") as f:
|
|
json.dump(
|
|
{"step": self.current_step, "width": self.width, "height": self.height},
|
|
f,
|
|
)
|
|
print(f"Checkpoint saved at step {self.current_step}")
|
|
|
|
def load_checkpoint(self, checkpoint):
|
|
"""Load checkpoint - called automatically by Tune during restoration"""
|
|
checkpoint_path = os.path.join(checkpoint, "checkpoint.json")
|
|
with open(checkpoint_path, "r") as f:
|
|
state = json.load(f)
|
|
self.current_step = state["step"]
|
|
self.width = state["width"]
|
|
self.height = state["height"]
|
|
print(f"Checkpoint loaded from step {self.current_step}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--smoke-test", action="store_true", help="Finish quickly for testing"
|
|
)
|
|
args, _ = parser.parse_known_args()
|
|
|
|
print(
|
|
"=" * 60,
|
|
"Ray Tune Example: Smart Checkpointing with custom SHOULD_CHECKPOINT key",
|
|
"=" * 60,
|
|
"",
|
|
"This example demonstrates how to set the SHOULD_CHECKPOINT key in a callback",
|
|
"to implement intelligent checkpointing based on training progress.",
|
|
"",
|
|
"Key features:",
|
|
"- Callback-driven checkpointing by setting result[SHOULD_CHECKPOINT] = True",
|
|
"- Checkpoints triggered by performance improvements",
|
|
"- Milestone-based checkpointing every 10 steps",
|
|
"- Instability detection (high variance in recent losses)",
|
|
"- Automatic checkpoint save/load via class trainable",
|
|
sep="\n",
|
|
)
|
|
|
|
# Create the smart checkpoint callback
|
|
checkpoint_callback = SmartCheckpointCallback(
|
|
checkpoint_on_improvement=True, checkpoint_on_instability=True
|
|
)
|
|
|
|
tuner = tune.Tuner(
|
|
OptimizationTrainable,
|
|
run_config=tune.RunConfig(
|
|
name="smart_checkpoint_test",
|
|
stop={"training_iteration": 1 if args.smoke_test else 20},
|
|
callbacks=[checkpoint_callback], # Add our custom callback
|
|
# Disable automatic periodic checkpointing to show callback control
|
|
checkpoint_config=tune.CheckpointConfig(
|
|
checkpoint_frequency=0, # Disable periodic checkpointing
|
|
checkpoint_at_end=True, # Still checkpoint at the end
|
|
),
|
|
),
|
|
tune_config=tune.TuneConfig(
|
|
metric="mean_loss",
|
|
mode="min",
|
|
num_samples=3,
|
|
),
|
|
param_space={
|
|
"width": tune.randint(10, 100),
|
|
"height": tune.loguniform(10, 100),
|
|
},
|
|
)
|
|
|
|
print(
|
|
"Starting hyperparameter tuning with smart checkpointing...",
|
|
"Watch for checkpoint messages triggered by the callback!",
|
|
sep="\n",
|
|
)
|
|
|
|
results = tuner.fit()
|
|
best_result = results.get_best_result()
|
|
print(
|
|
"\n" + "=" * 60,
|
|
"RESULTS",
|
|
"=" * 60,
|
|
f"Best hyperparameters: {best_result.config}",
|
|
f"Best checkpoint: {best_result.checkpoint}",
|
|
"",
|
|
"The checkpoints were triggered by the SmartCheckpointCallback",
|
|
sep="\n",
|
|
)
|