使用 EvoX 进行高效 HPO

使用 EvoX 进行高效 HPO

在本章中,我们将探讨如何使用 EvoX 进行超参数优化(HPO)。

HPO 在许多机器学习任务中起着至关重要的作用,但由于其高昂的计算成本(有时需要数天才能处理完)以及部署方面的挑战,往往被忽视。

通过 EvoX,我们可以使用 HPOProblemWrapper 简化 HPO 部署,并利用 vmap 方法和 GPU 加速实现高效计算。

将工作流转换为问题

HPO 结构

使用 EvoX 部署 HPO 的关键在于使用 HPOProblemWrapperworkflows 转换为 problems。转换完成后,我们可以将 workflows 视为标准的 problems。‘HPO 问题’的输入由超参数组成,输出则是评估指标。

核心组件 —— HPOProblemWrapper

为了确保 HPOProblemWrapper 能够识别超参数,我们需要使用 Parameter 对其进行封装。通过这一简单的步骤,超参数将被自动识别。

class ExampleAlgorithm(Algorithm):
    def __init__(self,...): 
        self.omega = Parameter([1.0, 2.0]) # wrap the hyper-parameters with `Parameter`
        self.beta = Parameter(0.1)
        pass

    def step(self):
        # run algorithm step depending on the value of self.omega and self.beta
        pass

使用 HPOFitnessMonitor

我们提供了一个 HPOFitnessMonitor,支持计算多目标问题的 ‘IGD’ 和 ‘HV’ 指标,以及单目标问题的最小值。

需要注意的是,HPOFitnessMonitor 是一个为 HPO 问题设计的基础监控器。您也可以参考 使用自定义算法部署 HPO 中概述的方法,灵活地创建自定义监控器。

一个简单的例子

在这里,我们将演示一个使用 EvoX 进行 HPO 的简单示例。具体来说,我们将使用 PSO 算法来优化用于求解 Sphere 问题的 PSO 算法的超参数。

请注意,本章仅提供 HPO 部署的简要概述。有关更详细的指南,请参阅 使用自定义算法部署 HPO

首先,让我们导入必要的模块。

import torch

from evox.algorithms.pso_variants.pso import PSO
from evox.core import Problem
from evox.problems.hpo_wrapper import HPOFitnessMonitor, HPOProblemWrapper
from evox.workflows import EvalMonitor, StdWorkflow

接下来,我们定义一个简单的 Sphere 问题。

class Sphere(Problem):
    def __init__(self):
        super().__init__()

    def evaluate(self, x: torch.Tensor):
        return (x * x).sum(-1)

接下来,我们可以使用 StdWorkflow 来封装 problemalgorithmmonitor。然后,我们使用 HPOProblemWrapperStdWorkflow 转换为 HPO 问题。

# the inner loop is a PSO algorithm with a population size of 50
torch.set_default_device("cuda" if torch.cuda.is_available() else "cpu")
inner_algo = PSO(50, -10 * torch.ones(10), 10 * torch.ones(10))
inner_prob = Sphere()
inner_monitor = HPOFitnessMonitor()
inner_workflow = StdWorkflow(inner_algo, inner_prob, monitor=inner_monitor)
# Transform the inner workflow to an HPO problem
hpo_prob = HPOProblemWrapper(iterations=30, num_instances=128, workflow=inner_workflow, copy_init_state=True)

HPOProblemWrapper 接受 4 个参数:

  1. iterations:优化过程中执行的迭代次数。
  2. num_instances:优化过程中并行执行的实例数量。
  3. workflow:优化过程中使用的工作流。
  4. copy_init_state:是否为每次评估复制工作流的初始状态。默认为 True。如果您的工作流包含对初始状态中的张量进行原位(IN-PLACE)修改的操作,则应将其设置为 True。否则,您可以将其设置为 False 以节省内存。

我们可以验证 HPOProblemWrapper 是否正确识别了我们定义的超参数。由于 5 个实例的超参数未做任何修改,因此所有实例的超参数应保持一致。

params = hpo_prob.get_init_params()
print("init params:\n", params)

我们还可以定义一组自定义的超参数值。务必确保超参数组的数量与 HPOProblemWrapper 中的实例数量相匹配。此外,自定义超参数必须以字典形式提供,且其值需使用 Parameter 进行封装。

params = hpo_prob.get_init_params()
# since we have 128 instances, we need to pass 128 sets of hyperparameters
params["algorithm.w"] = torch.nn.Parameter(torch.rand(128, 1), requires_grad=False)
params["algorithm.phi_p"] = torch.nn.Parameter(torch.rand(128, 1), requires_grad=False)
params["algorithm.phi_g"] = torch.nn.Parameter(torch.rand(128, 1), requires_grad=False)
result = hpo_prob.evaluate(params)
print("The result of the first 3 parameter sets:\n", result[:3])

现在,我们使用 PSO 算法来优化 PSO 算法的超参数。

务必确保 PSO 的种群大小与实例数量相匹配,否则可能会发生意外错误。

此外,解需要在外部工作流中进行转换,因为 HPOProblemWrapper 要求输入必须是字典形式。

class solution_transform(torch.nn.Module):
    def forward(self, x: torch.Tensor):
        return {
            "algorithm.w": x[:, 0],
            "algorithm.phi_p": x[:, 1],
            "algorithm.phi_g": x[:, 2],
        }


outer_algo = PSO(128, 0 * torch.ones(3), 3 * torch.ones(3))  # search each hyperparameter in the range [0, 3]
monitor = EvalMonitor(full_sol_history=False)
outer_workflow = StdWorkflow(outer_algo, hpo_prob, monitor=monitor, solution_transform=solution_transform())
outer_workflow.init_step()
compiled_step = torch.compile(outer_workflow.step)
for _ in range(100):
    compiled_step()
monitor = outer_workflow.get_submodule("monitor")
print("params:\n", monitor.topk_solutions, "\n")
print("result:\n", monitor.topk_fitness)
monitor.plot()