wehub-resource-sync
137 lines
3.4 KiB
Python
137 lines
3.4 KiB
Python
# flake8: noqa
|
|
# isort: skip_file
|
|
|
|
from pathlib import Path
|
|
import tempfile
|
|
|
|
import ray.train
|
|
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
|
|
|
|
|
|
def train_fn(config):
|
|
for i in range(3):
|
|
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
|
Path(temp_checkpoint_dir).joinpath("model.pt").touch()
|
|
ray.train.report(
|
|
{"loss": i},
|
|
checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
|
|
)
|
|
|
|
return {"total loss": 3}
|
|
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_fn, scaling_config=ray.train.ScalingConfig(num_workers=2)
|
|
)
|
|
|
|
|
|
# __run_config_start__
|
|
import os
|
|
|
|
from ray.train import RunConfig
|
|
|
|
run_config = RunConfig(
|
|
# Name of the training run (directory name).
|
|
name="my_train_run",
|
|
# The experiment results will be saved to: storage_path/name
|
|
storage_path=os.path.expanduser("~/ray_results"),
|
|
# storage_path="s3://my_bucket/tune_results",
|
|
)
|
|
# __run_config_end__
|
|
|
|
# __checkpoint_config_start__
|
|
from ray.train import RunConfig, CheckpointConfig
|
|
|
|
# Example 1: Only keep the 2 *most recent* checkpoints and delete the others.
|
|
run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=2))
|
|
|
|
|
|
# Example 2: Only keep the 2 *best* checkpoints and delete the others.
|
|
run_config = RunConfig(
|
|
checkpoint_config=CheckpointConfig(
|
|
num_to_keep=2,
|
|
# *Best* checkpoints are determined by these params:
|
|
checkpoint_score_attribute="mean_accuracy",
|
|
checkpoint_score_order="max",
|
|
),
|
|
# This will store checkpoints on S3.
|
|
storage_path="s3://remote-bucket/location",
|
|
)
|
|
# __checkpoint_config_end__
|
|
|
|
|
|
# __result_metrics_start__
|
|
result = trainer.fit()
|
|
|
|
print("Observed metrics:", result.metrics)
|
|
# __result_metrics_end__
|
|
|
|
|
|
# __result_dataframe_start__
|
|
df = result.metrics_dataframe
|
|
print("Minimum loss", min(df["loss"]))
|
|
# __result_dataframe_end__
|
|
|
|
|
|
# __result_return_value_start__
|
|
print("Returned data", result.return_value)
|
|
# __result_return_value_end__
|
|
|
|
|
|
# __result_checkpoint_start__
|
|
print("Last checkpoint:", result.checkpoint)
|
|
|
|
with result.checkpoint.as_directory() as tmpdir:
|
|
# Load model from directory
|
|
...
|
|
# __result_checkpoint_end__
|
|
|
|
# __result_best_checkpoint_start__
|
|
# Print available checkpoints
|
|
for checkpoint, metrics in result.best_checkpoints:
|
|
print("Loss", metrics["loss"], "checkpoint", checkpoint)
|
|
|
|
# Get checkpoint with minimal loss
|
|
best_checkpoint = min(
|
|
result.best_checkpoints, key=lambda checkpoint: checkpoint[1]["loss"]
|
|
)[0]
|
|
|
|
with best_checkpoint.as_directory() as tmpdir:
|
|
# Load model from directory
|
|
...
|
|
# __result_best_checkpoint_end__
|
|
|
|
import pyarrow
|
|
|
|
# __result_path_start__
|
|
result_path: str = result.path
|
|
result_filesystem: pyarrow.fs.FileSystem = result.filesystem
|
|
|
|
print(f"Results location (fs, path) = ({result_filesystem}, {result_path})")
|
|
# __result_path_end__
|
|
|
|
|
|
# __result_restore_start__
|
|
from ray.train import Result
|
|
|
|
restored_result = Result.from_path(result_path)
|
|
print("Restored loss", restored_result.metrics["loss"])
|
|
# __result_restore_end__
|
|
|
|
|
|
def error_train_fn(config):
|
|
raise RuntimeError("Simulated training error")
|
|
|
|
|
|
trainer = DataParallelTrainer(
|
|
error_train_fn, scaling_config=ray.train.ScalingConfig(num_workers=1)
|
|
)
|
|
|
|
# __result_error_start__
|
|
try:
|
|
result = trainer.fit()
|
|
except ray.train.TrainingFailedError as e:
|
|
if isinstance(e, ray.train.WorkerGroupError):
|
|
print(e.worker_failures)
|
|
# __result_error_end__
|