如何为 Ray Tune 实验定义停止条件#

在运行 Tune 实验时,提前确定理想的训练时长可能具有挑战性。Tune 中的停止条件可以根据特定条件终止训练。

例如,可以设置实验在以下情况下停止

  1. 设置实验,使其在 N 个 epoch 后停止,或当报告的评估分数超过特定阈值时停止,以先到者为准。

  2. T 秒后停止实验。

  3. 当 trials 遇到运行时错误时终止。

  4. 利用 Tune 的提前停止调度器,尽早停止表现不佳的 trials。

本用户指南将说明如何在 Tune 实验中实现这些类型的停止条件。

对于所有代码示例,我们使用以下训练函数进行演示

from ray import tune
import time

def my_trainable(config):
    i = 1
    while True:
        # Do some training...
        time.sleep(1)

        # Report some metrics for demonstration...
        tune.report({"mean_accuracy": min(i / 10, 1.0)})
        i += 1

手动停止 Tune 实验#

如果您向运行 Tuner.fit() 的进程发送 SIGINT 信号(通常是您在终端按下 Ctrl+C 时发生的情况),Ray Tune 会优雅地关闭训练并保存最终的实验状态。

注意

强制终止 Tune 实验,例如通过多次按下 Ctrl+C 命令,将不会给 Tune 保存最后一次实验状态快照的机会。如果您将来恢复实验,这可能会导致使用过时的状态恢复。

Ray Tune 还接受 SIGUSR1 信号来优雅地中断训练。当在远程 Ray 任务中运行 Ray Tune 时,应使用此信号,因为 Ray 默认会过滤掉 SIGINTSIGTERM 信号。

使用基于指标的条件停止#

除了手动停止,Tune 还提供了多种以编程方式停止实验的方法。最简单的方法是使用基于指标的条件。这些是确定实验何时应停止的固定阈值集。

您可以使用字典、函数或自定义 Stopper 来实现停止条件。

如果传入字典,键可以是 Function API 中 tune.report 的返回值或 Class API 中 step() 的任何字段。

注意

这包括 自动填充的指标,例如 training_iteration

在下面的示例中,每个 trial 将在完成 10 次迭代后停止,或者当达到 0.8 或更高的平均准确率时停止。

这些指标假定为**递增**的,因此一旦报告的指标超过字典中指定的阈值,trial 就会停止。

from ray import tune

tuner = tune.Tuner(
    my_trainable,
    run_config=tune.RunConfig(stop={"training_iteration": 10, "mean_accuracy": 0.8}),
)
result_grid = tuner.fit()

为了获得更大的灵活性,您可以传入一个函数。如果传入一个函数,它必须接受 (trial_id: str, result: dict) 作为参数,并返回一个布尔值(如果 trial 应停止则返回 True,否则返回 False)。

在下面的示例中,每个 trial 将在完成 10 次迭代后停止,或者当达到 0.8 或更高的平均准确率时停止。

from ray import tune


def stop_fn(trial_id: str, result: dict) -> bool:
    return result["mean_accuracy"] >= 0.8 or result["training_iteration"] >= 10


tuner = tune.Tuner(my_trainable, run_config=tune.RunConfig(stop=stop_fn))
result_grid = tuner.fit()

最后,您可以实现 Stopper 接口,以根据自定义停止条件停止单个 trial 甚至整个实验。例如,以下示例在任何单个 trial 达到标准后停止所有 trial,并阻止新的 trial 开始。

from ray import tune
from ray.tune import Stopper


class CustomStopper(Stopper):
    def __init__(self):
        self.should_stop = False

    def __call__(self, trial_id: str, result: dict) -> bool:
        if not self.should_stop and result["mean_accuracy"] >= 0.8:
            self.should_stop = True
        return self.should_stop

    def stop_all(self) -> bool:
        """Returns whether to stop trials and prevent new ones from starting."""
        return self.should_stop


stopper = CustomStopper()
tuner = tune.Tuner(
    my_trainable,
    run_config=tune.RunConfig(stop=stopper),
    tune_config=tune.TuneConfig(num_samples=2),
)
result_grid = tuner.fit()

在示例中,一旦任何 trial 达到 0.8 或更高的 mean_accuracy,所有 trial 将停止。

注意

stop_all 返回 True 时,当前正在运行的 trial 不会立即停止。它们将在完成当前的训练迭代(在 tune.reportstep 之后)停止。

Ray Tune 提供了一组开箱即用的 stopper 类。请参阅 Stopper 文档。

在一定时间后停止 trials#

有两种选择可以基于时间停止 Tune 实验:在指定超时后单独停止 trials,或在一定时间后停止整个实验。

通过超时单独停止 trials#

您可以使用上面描述的字典停止条件,使用 Tune 自动填充的 time_total_s 指标。

from ray import tune

tuner = tune.Tuner(
    my_trainable,
    # Stop a trial after it's run for more than 5 seconds.
    run_config=tune.RunConfig(stop={"time_total_s": 5}),
)
result_grid = tuner.fit()

注意

如果您使用 Function Trainable API,则需要通过 tune.report 包含一些中间报告。每次报告都会自动记录 trial 的 time_total_s,这使得 Tune 可以根据时间作为指标进行停止。

如果训练循环在某个地方挂起,Tune 将无法拦截训练并为您停止 trial。在这种情况下,您可以在训练循环中显式实现超时逻辑。

通过超时停止实验#

使用 TuneConfig(time_budget_s) 配置,告诉 Tune 在 time_budget_s 秒后停止实验。

from ray import tune

# Stop the entire experiment after ANY trial has run for more than 5 seconds.
tuner = tune.Tuner(my_trainable, tune_config=tune.TuneConfig(time_budget_s=5.0))
result_grid = tuner.fit()

注意

如果您使用 Function Trainable API,则需要通过 tune.report 包含一些中间报告,原因与上面相同。

因 trial 失败而停止#

除了根据 trial 的表现停止它们,您还可以让整个实验在任何 trial 遇到运行时错误时停止。为此,您可以使用 ray.tune.FailureConfig 类。

有了此配置,如果任何 trial 遇到错误,整个实验将立即停止。

from ray import tune
import time


def my_failing_trainable(config):
    if config["should_fail"]:
        raise RuntimeError("Failing (on purpose)!")
    # Do some training...
    time.sleep(10)
    tune.report({"mean_accuracy": 0.9})


tuner = tune.Tuner(
    my_failing_trainable,
    param_space={"should_fail": tune.grid_search([True, False])},
    run_config=tune.RunConfig(failure_config=tune.FailureConfig(fail_fast=True)),
)
result_grid = tuner.fit()

这对于调试具有大量 trial 的 Tune 实验非常有用。

使用 Tune 调度器进行提前停止#

停止 Tune 实验的另一种方法是使用提前停止调度器。这些调度器会监控 trial 的表现,并在它们没有取得足够进展时尽早停止它们。

AsyncHyperBandSchedulerHyperBandForBOHB 是 Tune 中内置的提前停止调度器的示例。有关完整列表以及更实际的示例,请参阅 Tune 调度器 API 参考

在以下示例中,我们同时使用了字典停止条件和提前停止条件。

from ray import tune
from ray.tune.schedulers import AsyncHyperBandScheduler


scheduler = AsyncHyperBandScheduler(time_attr="training_iteration")

tuner = tune.Tuner(
    my_trainable,
    run_config=tune.RunConfig(stop={"training_iteration": 10}),
    tune_config=tune.TuneConfig(
        scheduler=scheduler, num_samples=2, metric="mean_accuracy", mode="max"
    ),
)
result_grid = tuner.fit()

总结#

在本用户指南中,我们学习了如何使用指标、trial 错误和提前停止调度器来停止 Tune 实验。

有关更多信息,请参阅以下资源: