wehub-resource-sync
71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
# If want to use checkpointing with a custom training function (not a Ray
|
|
# integration like PyTorch or Tensorflow), your function can read/write
|
|
# checkpoint through the ``ray.tune.report(metrics, checkpoint=...)`` API.
|
|
import argparse
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import time
|
|
|
|
from ray import tune
|
|
from ray.tune import Checkpoint
|
|
|
|
|
|
def evaluation_fn(step, width, height):
|
|
time.sleep(0.1)
|
|
return (0.1 + width * step / 100) ** (-1) + height * 0.1
|
|
|
|
|
|
def train_func(config):
|
|
step = 0
|
|
width, height = config["width"], config["height"]
|
|
|
|
checkpoint = tune.get_checkpoint()
|
|
if checkpoint:
|
|
with checkpoint.as_directory() as checkpoint_dir:
|
|
with open(os.path.join(checkpoint_dir, "checkpoint.json")) as f:
|
|
state = json.load(f)
|
|
step = state["step"] + 1
|
|
|
|
for current_step in range(step, 100):
|
|
intermediate_score = evaluation_fn(current_step, width, height)
|
|
|
|
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
|
with open(os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w") as f:
|
|
json.dump({"step": current_step}, f)
|
|
tune.report(
|
|
{"iterations": current_step, "mean_loss": intermediate_score},
|
|
checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--smoke-test", action="store_true", help="Finish quickly for testing"
|
|
)
|
|
args, _ = parser.parse_known_args()
|
|
|
|
tuner = tune.Tuner(
|
|
train_func,
|
|
run_config=tune.RunConfig(
|
|
name="hyperband_test",
|
|
stop={"training_iteration": 1 if args.smoke_test else 10},
|
|
),
|
|
tune_config=tune.TuneConfig(
|
|
metric="mean_loss",
|
|
mode="min",
|
|
num_samples=5,
|
|
),
|
|
param_space={
|
|
"steps": 10,
|
|
"width": tune.randint(10, 100),
|
|
"height": tune.loguniform(10, 100),
|
|
},
|
|
)
|
|
results = tuner.fit()
|
|
best_result = results.get_best_result()
|
|
print("Best hyperparameters: ", best_result.config)
|
|
best_checkpoint = best_result.checkpoint
|
|
print("Best checkpoint: ", best_checkpoint)
|