# If want to use checkpointing with a custom training function (not a Ray
# integration like PyTorch or Tensorflow), you must expose a
# ``checkpoint_dir`` argument in the function signature, and call
# ``tune.checkpoint_dir``:
import os
import time
import json
import argparse
from ray import tune
def evaluation_fn(step, width, height):
time.sleep(0.1)
return (0.1 + width * step / 100) ** (-1) + height * 0.1
def train_func(config, checkpoint_dir=None):
start = 0
width, height = config["width"], config["height"]
if checkpoint_dir:
with open(os.path.join(checkpoint_dir, "checkpoint")) as f:
state = json.loads(f.read())
start = state["step"] + 1
for step in range(start, 100):
intermediate_score = evaluation_fn(step, width, height)
# Obtain a checkpoint directory
with tune.checkpoint_dir(step=step) as checkpoint_dir:
path = os.path.join(checkpoint_dir, "checkpoint")
with open(path, "w") as f:
f.write(json.dumps({"step": step}))
tune.report(iterations=step, mean_loss=intermediate_score)
# You can restore a single trial checkpoint by using
# ``tune.run(restore=<checkpoint_dir>)`` By doing this, you can change
# whatever experiments' configuration such as the experiment's name.
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--smoke-test", action="store_true", help="Finish quickly for testing"
)
parser.add_argument(
"--server-address",
type=str,
default=None,
required=False,
help="The address of server to connect to if using Ray Client.",
)
args, _ = parser.parse_known_args()
if args.server_address:
import ray
ray.init(f"ray://{args.server_address}")
analysis = tune.run(
train_func,
name="hyperband_test",
metric="mean_loss",
mode="min",
num_samples=5,
stop={"training_iteration": 1 if args.smoke_test else 10},
config={
"steps": 10,
"width": tune.randint(10, 100),
"height": tune.loguniform(10, 100),
},
)
print("Best hyperparameters: ", analysis.best_config)
print("Best checkpoint directory: ", analysis.best_checkpoint)
with open(os.path.join(analysis.best_checkpoint, "checkpoint"), "r") as f:
print("Best checkpoint: ", json.load(f))
