Reproducibility and Checkpointing#
Makes long-running optimizations reproducible and lets you resume them partway through.
Checkpointing uses the low-level Optimizer API.
Construct Optimizer(problem), configure components by chaining set_*(), and run it with run().
Resume from a saved state with run_from().
What you’ll be able to do
By the end of this page, you’ll be able to ensure reproducibility with a random seed, save and load checkpoints, and resume an optimization.
For a single run with the high-level API, use the high-level API; for finer control of run state, see the low-level API. This page shows how to save a reproducible low-level run and resume it partway through.
Reproduce runs with a random seed#
Passing the same seed to Optimizer(problem, seed=...) initializes all random-number-using processes in the same sequence, producing identical results.
import numpy as np
from saealib import (
Problem,
Optimizer,
GA,
CrossoverBLXAlpha,
MutationUniform,
SequentialSelection,
TruncationSelection,
RBFSurrogate,
GaussianKernel,
LocalSurrogateManager,
MeanPrediction,
IndividualBasedStrategy,
LHSInitializer,
Termination,
max_fe,
)
def expensive_func(x):
return np.sum(x**2)
DIM = 10
SEED = 0
problem = Problem(
func=expensive_func,
dim=DIM,
n_obj=1,
direction=np.array([-1.0]),
lb=[-5.0] * DIM,
ub=[5.0] * DIM,
)
def build_optimizer(max_fe_value):
return (
Optimizer(problem, seed=SEED)
.set_initializer(
LHSInitializer(n_init_archive=5 * DIM, n_init_population=4 * DIM, seed=SEED)
)
.set_algorithm(
GA(
crossover=CrossoverBLXAlpha(0.7, 0.4),
mutation=MutationUniform(0.3),
parent_selection=SequentialSelection(),
survivor_selection=TruncationSelection(),
)
)
.set_surrogate_manager(
LocalSurrogateManager(RBFSurrogate(kernel=GaussianKernel()))
)
.set_acquisition(MeanPrediction())
.set_strategy(IndividualBasedStrategy(evaluation_ratio=0.1))
.set_termination(Termination(max_fe(max_fe_value)))
)
ctx1 = build_optimizer(300).run()
ctx2 = build_optimizer(300).run()
print(np.allclose(ctx1.archive.get_array("f"), ctx2.archive.get_array("f"))) # True
build_optimizer is used in the following sections too, to rebuild an Optimizer with the same component configuration.
To preserve reproducibility, recreate the SurrogateManager and AcquisitionFunction with the same configuration.
The current API does not support passing an AcquisitionFunction as an argument to LocalSurrogateManager.
Save and resume a checkpoint#
The ctx returned by run() can be saved to a single npz file with ctx.save(path).
ctx = build_optimizer(200).run()
ctx.save("checkpoint.npz")
A saved checkpoint can be loaded with OptimizationState.load(path, problem), and passing it to Optimizer.run_from(ctx) resumes from where it left off.
from saealib import OptimizationState
loaded_ctx = OptimizationState.load("checkpoint.npz", problem)
resumed_ctx = build_optimizer(600).run_from(loaded_ctx)
print(resumed_ctx.fe) # includes the evaluations from before saving
print(resumed_ctx.data["resumed"]) # True
ctx.data["resumed"] is a flag set to True only on a context resumed via run_from().
From a callback such as RunStartEvent, it can be accessed as event.ctx.data["resumed"].
Save checkpoints automatically#
Passing checkpoint_path to run()/iterate() saves automatically every checkpoint_interval generations.
checkpoint_path is treated as a directory rather than a single file, and per-generation snapshots are created, named checkpoint_{gen:06d}.npz.
ctx = build_optimizer(300).run(checkpoint_path="checkpoints", checkpoint_interval=5)
To resume, load the most recent snapshot in the directory.
from pathlib import Path
latest = sorted(Path("checkpoints").glob("checkpoint_*.npz"))[-1]
loaded_ctx = OptimizationState.load(latest, problem)
If you don’t want to leave snapshots behind after a successful run, specify checkpoint_delete_on_success=True (the directory itself is kept; only the files inside it are deleted).
ctx = build_optimizer(300).run(
checkpoint_path="checkpoints",
checkpoint_interval=5,
checkpoint_delete_on_success=True,
)
Save in pickle format#
npz saves only ctx, but pickle format can save the entire Optimizer, including the fitted surrogate parameters.
optimizer = build_optimizer(200)
ctx = optimizer.run()
optimizer.save_pickle(ctx, "checkpoint.pkl")
loaded_optimizer, loaded_ctx = Optimizer.load_pickle("checkpoint.pkl")
A UserWarning about the Python or library version may appear at runtime.
An Optimizer containing objects that the standard pickle cannot serialize — such as a lambda used in Termination — cannot be pickle-saved.
Use CheckpointCallback directly#
The checkpoint_path argument of run() simply registers a CheckpointCallback internally.
To wire up the same behavior explicitly, register a CheckpointCallback on cbmanager.
from saealib import CheckpointCallback
optimizer = build_optimizer(300)
callback = CheckpointCallback("checkpoints", interval=5, optimizer=optimizer)
callback.register(optimizer.cbmanager)
ctx = optimizer.run()
The optimizer argument is required when specifying format="pickle" or format="both".