chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,221 @@
|
||||
# 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",
|
||||
)
|
||||
Reference in New Issue
Block a user