wehub-resource-sync
312 lines
11 KiB
Python
312 lines
11 KiB
Python
# __validation_fn_simple_start__
|
|
import os
|
|
import torch
|
|
|
|
import ray.train
|
|
import ray.data
|
|
|
|
# Define Ray Data validation dataset outside validation function because it is not json serializable
|
|
validation_dataset = ...
|
|
|
|
|
|
def validation_fn(checkpoint: ray.train.Checkpoint) -> dict:
|
|
# Load the checkpoint
|
|
model = ...
|
|
with checkpoint.as_directory() as checkpoint_dir:
|
|
model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
|
|
model.load_state_dict(model_state_dict)
|
|
model.eval()
|
|
|
|
# Perform validation on the data
|
|
total_accuracy = 0
|
|
with torch.no_grad():
|
|
for batch in validation_dataset.iter_torch_batches(batch_size=128):
|
|
images, labels = batch["image"], batch["label"]
|
|
outputs = model(images)
|
|
total_accuracy += (outputs.argmax(1) == labels).sum().item()
|
|
return {"score": total_accuracy / len(validation_dataset)}
|
|
# __validation_fn_simple_end__
|
|
|
|
# __validation_fn_torch_trainer_start__
|
|
import torchmetrics
|
|
from torch.nn import CrossEntropyLoss
|
|
|
|
import ray.train.torch
|
|
from ray.data import ExecutionOptions
|
|
|
|
|
|
def eval_only_train_fn(config_dict: dict) -> dict:
|
|
# Load the checkpoint
|
|
model = ...
|
|
with config_dict["checkpoint"].as_directory() as checkpoint_dir:
|
|
model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
|
|
model.load_state_dict(model_state_dict)
|
|
model.cuda().eval()
|
|
|
|
# Set up metrics and data loaders
|
|
criterion = CrossEntropyLoss()
|
|
mean_valid_loss = torchmetrics.MeanMetric().cuda()
|
|
test_data_shard = ray.train.get_dataset_shard("validation")
|
|
test_dataloader = test_data_shard.iter_torch_batches(batch_size=128)
|
|
|
|
# Compute metric and return it directly from the train function
|
|
with torch.no_grad():
|
|
for batch in test_dataloader:
|
|
images, labels = batch["image"], batch["label"]
|
|
outputs = model(images)
|
|
loss = criterion(outputs, labels)
|
|
mean_valid_loss(loss)
|
|
return {"score": mean_valid_loss.compute().item()}
|
|
|
|
|
|
def validation_fn(checkpoint: ray.train.Checkpoint, train_run_name: str, epoch: int) -> dict:
|
|
trainer = ray.train.torch.TorchTrainer(
|
|
eval_only_train_fn,
|
|
train_loop_config={"checkpoint": checkpoint},
|
|
scaling_config=ray.train.ScalingConfig(
|
|
num_workers=2, use_gpu=True, accelerator_type="A10G"
|
|
),
|
|
# Give unique name to validation run so it does not attempt to load placeholder checkpoint.
|
|
# Also allows you to better associate training runs with validation runs.
|
|
run_config=ray.train.RunConfig(
|
|
name=f"{train_run_name}_validation_epoch_{epoch}"
|
|
),
|
|
# Use weaker GPUs for validation
|
|
datasets={"validation": validation_dataset},
|
|
# Pin to the "validation" subcluster so it doesn't compete with
|
|
# training. See https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
|
|
dataset_config=ray.train.DataConfig(
|
|
execution_options={
|
|
"validation": ExecutionOptions(
|
|
label_selector={"ray-subcluster": "validation"}
|
|
),
|
|
},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
# return_value holds the value returned by train function of worker 0
|
|
return result.return_value
|
|
# __validation_fn_torch_trainer_end__
|
|
|
|
# __validation_fn_map_batches_start__
|
|
import ray.data
|
|
|
|
|
|
class Predictor:
|
|
def __init__(self, checkpoint: ray.train.Checkpoint):
|
|
self.model = ...
|
|
with checkpoint.as_directory() as checkpoint_dir:
|
|
model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
|
|
self.model.load_state_dict(model_state_dict)
|
|
self.model.cuda().eval()
|
|
|
|
def __call__(self, batch: dict) -> dict:
|
|
image = torch.as_tensor(batch["image"], dtype=torch.float32, device="cuda")
|
|
label = torch.as_tensor(batch["label"], dtype=torch.float32, device="cuda")
|
|
pred = self.model(image)
|
|
return {"res": (pred.argmax(1) == label).cpu().numpy()}
|
|
|
|
|
|
# Construct ``validation_dataset`` under a DataContext copy pinned to the
|
|
# "validation" subcluster. ``Dataset.context`` is a deep copy of the
|
|
# current context taken at construction, so the selector is baked in and
|
|
# every downstream operator (including the ``map_batches`` below) inherits
|
|
# it — no in-function mutation needed. See
|
|
# https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
|
|
ctx = ray.data.DataContext.get_current().copy()
|
|
ctx.execution_options.label_selector = {"ray-subcluster": "validation"}
|
|
with ray.data.DataContext.current(ctx):
|
|
validation_dataset = ray.data.read_parquet(...)
|
|
|
|
|
|
def validation_fn(checkpoint: ray.train.Checkpoint) -> dict:
|
|
# Set name to avoid confusion; default name is "Dataset"
|
|
validation_dataset.set_name("validation")
|
|
eval_res = validation_dataset.map_batches(
|
|
Predictor,
|
|
batch_size=128,
|
|
num_gpus=1,
|
|
fn_constructor_kwargs={"checkpoint": checkpoint},
|
|
concurrency=2,
|
|
)
|
|
mean = eval_res.mean(["res"])
|
|
return {
|
|
"score": mean,
|
|
}
|
|
# __validation_fn_map_batches_end__
|
|
|
|
# __validation_fn_report_start__
|
|
import tempfile
|
|
|
|
from ray.data import ExecutionOptions
|
|
from ray.train import ValidationConfig, ValidationTaskConfig
|
|
|
|
|
|
def train_func(config: dict) -> None:
|
|
...
|
|
epochs = ...
|
|
model = ...
|
|
rank = ray.train.get_context().get_world_rank()
|
|
for epoch in epochs:
|
|
... # training step
|
|
if rank == 0:
|
|
training_metrics = {"loss": ..., "epoch": epoch}
|
|
local_checkpoint_dir = tempfile.mkdtemp()
|
|
torch.save(
|
|
model.module.state_dict(),
|
|
os.path.join(local_checkpoint_dir, "model.pt"),
|
|
)
|
|
ray.train.report(
|
|
training_metrics,
|
|
checkpoint=ray.train.Checkpoint.from_directory(local_checkpoint_dir),
|
|
checkpoint_upload_mode=ray.train.CheckpointUploadMode.ASYNC,
|
|
validation=ValidationTaskConfig(fn_kwargs={
|
|
"train_run_name": ray.train.get_context().get_experiment_name(),
|
|
"epoch": epoch,
|
|
}),
|
|
)
|
|
else:
|
|
ray.train.report({}, None)
|
|
|
|
|
|
def run_trainer() -> ray.train.Result:
|
|
# 1) Construction-time tasks (parquet schema inference, file listing)
|
|
# read the current DataContext. Pin them to "training" with a copy of
|
|
# the DataContext applied via the DataContext.current() context
|
|
# manager — scoped to the `with` block so it doesn't leak. See
|
|
# https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
|
|
ctx = ray.data.DataContext.get_current().copy()
|
|
ctx.execution_options.label_selector = {"ray-subcluster": "training"}
|
|
with ray.data.DataContext.current(ctx):
|
|
train_dataset = ray.data.read_parquet(...)
|
|
|
|
trainer = ray.train.torch.TorchTrainer(
|
|
train_func,
|
|
validation_config=ValidationConfig(fn=validation_fn),
|
|
# Pass training dataset in datasets arg to split it across training workers
|
|
datasets={"train": train_dataset},
|
|
# 2) DataConfig.execution_options REPLACES ds.context.execution_options
|
|
# wholesale at training start, dropping anything not re-specified
|
|
# (including label_selector). Restate the selector here so per-worker
|
|
# ingest stays pinned to "training".
|
|
dataset_config=ray.train.DataConfig(
|
|
datasets_to_split=["train"],
|
|
execution_options={
|
|
"train": ExecutionOptions(
|
|
label_selector={"ray-subcluster": "training"}
|
|
),
|
|
},
|
|
),
|
|
scaling_config=ray.train.ScalingConfig(
|
|
num_workers=2,
|
|
use_gpu=True,
|
|
# Use powerful GPUs for training
|
|
accelerator_type="A100",
|
|
),
|
|
)
|
|
return trainer.fit()
|
|
# __validation_fn_report_end__
|
|
|
|
# __exp_tracking_same_run_wandb_start__
|
|
import wandb
|
|
import ray.train
|
|
from ray.train import ValidationConfig, ValidationTaskConfig
|
|
|
|
|
|
entity = "my_entity"
|
|
project = "my_project"
|
|
num_epochs = ...
|
|
|
|
|
|
def validation_fn(checkpoint: ray.train.Checkpoint, wandb_run_id: str, val_step: int) -> dict:
|
|
wandb.init(
|
|
entity=entity,
|
|
project=project,
|
|
settings=wandb.Settings(mode="shared", x_primary=False),
|
|
id=wandb_run_id,
|
|
)
|
|
score = ...
|
|
wandb.log({"validation/loss": score, "val_step": val_step})
|
|
wandb.finish() # flush the metrics
|
|
return {"validation/loss": score}
|
|
|
|
|
|
def train_func():
|
|
if ray.train.get_context().get_world_rank() == 0:
|
|
run = wandb.init(
|
|
entity=entity,
|
|
project=project,
|
|
settings=wandb.Settings(mode="shared", x_primary=True,)
|
|
)
|
|
wandb.define_metric("val_step", hidden=True)
|
|
wandb.define_metric("train_step", hidden=True)
|
|
wandb.define_metric("validation/loss", step_metric="val_step")
|
|
wandb.define_metric("train/loss", step_metric="train_step")
|
|
|
|
for epoch in range(num_epochs):
|
|
loss = ...
|
|
if ray.train.get_context().get_world_rank() == 0:
|
|
wandb.log({"train/loss": loss, "train_step": epoch})
|
|
checkpoint = ...
|
|
ray.train.report(
|
|
{"train/loss": loss},
|
|
checkpoint=checkpoint,
|
|
validation=ValidationTaskConfig(
|
|
fn_kwargs={"wandb_run_id": run.id, "val_step": epoch}
|
|
),
|
|
)
|
|
else:
|
|
ray.train.report({}, None)
|
|
|
|
if ray.train.get_context().get_world_rank() == 0:
|
|
wandb.finish()
|
|
|
|
|
|
# __exp_tracking_same_run_wandb_end__
|
|
|
|
# __exp_tracking_same_run_mlflow_start__
|
|
import mlflow
|
|
from mlflow.tracking import MlflowClient
|
|
import ray.train
|
|
from ray.train import ValidationConfig, ValidationTaskConfig
|
|
|
|
|
|
tracking_uri = "my_uri"
|
|
experiment_name = "my_experiment"
|
|
num_epochs = ...
|
|
|
|
def validation_fn(
|
|
checkpoint: ray.train.Checkpoint, mlflow_run_id: str, val_step: int
|
|
) -> dict:
|
|
client = MlflowClient(tracking_uri=tracking_uri)
|
|
score = ...
|
|
client.log_metric(mlflow_run_id, "val_score", score, step=val_step)
|
|
return {"val_score": score}
|
|
|
|
|
|
def train_func():
|
|
if ray.train.get_context().get_world_rank() == 0:
|
|
client = MlflowClient(tracking_uri=tracking_uri)
|
|
experiment = client.get_experiment_by_name(experiment_name)
|
|
run = client.create_run(experiment_id=experiment.experiment_id)
|
|
|
|
for epoch in range(num_epochs):
|
|
loss = ...
|
|
if ray.train.get_context().get_world_rank() == 0:
|
|
client.log_metric(run.info.run_id, "train_loss", loss, step=epoch)
|
|
checkpoint = ...
|
|
ray.train.report(
|
|
{"train_loss": loss},
|
|
checkpoint=checkpoint,
|
|
validation=ValidationTaskConfig(
|
|
fn_kwargs={"mlflow_run_id": run.info.run_id, "val_step": epoch}
|
|
),
|
|
)
|
|
else:
|
|
ray.train.report({}, None)
|
|
|
|
if ray.train.get_context().get_world_rank() == 0:
|
|
client.set_terminated(run.info.run_id)
|
|
|
|
# __exp_tracking_same_run_mlflow_end__
|