用Ax自定义一个BoTorch模型

0、前言

本篇译自BoTorch官方文档教程。

本教程中，我们将演示如何在Ax的SimpleExperimentAPI中，使用一个自定义的BoTorch模型。这使得我们可以利用Ax的便利性来运行贝叶斯优化的循环过程，同时保证建模过程的灵活性。

采集函数和优化采样的策略在许多相同的形式下（in much the same fashion）可以被去掉（can be swapped out）。该教程中提供了一个例子。

如果你不想按着模板来或者想动手实现建模优化的每一个步骤，参见添加链接描述此篇教程。

1、实现自定义模型

本教程中，我们用RBF核实现一个非常简单的Exact GP 模型（同分布噪声）。

模型的定义很直接，这里我们实现一个同时继承于ExactGP和GPyTorchModel的模型–这添加了botorch各种模型所需要的API调用。

注意：只要自定义模型遵循最小ModelAPI，botorch就允许实现这些自定义模型。更多资料参见该篇。

from botorch.models.gpytorch import GPyTorchModel
from gpytorch.distributions import MultivariateNormal
from gpytorch.means import ConstantMean
from gpytorch.models import ExactGP
from gpytorch.kernels import RBFKernel, ScaleKernel
from gpytorch.likelihoods import GaussianLikelihood
from gpytorch.mlls import ExactMarginalLogLikelihood
from gpytorch.priors import GammaPrior
from botorch.fit import fit_gpytorch_model
import random
import numpy as np
from ax import SearchSpace, RangeParameter, ParameterType, SimpleExperiment
from ax.modelbridge import get_sobol
from ax.modelbridge.factory import get_botorch


class SimpleCustomGP(ExactGP, GPyTorchModel):

    _num_outputs = 1  # to inform GPyTorchModel API
    
    def __init__(self, train_X, train_Y):
        # squeeze output dim before passing train_Y to ExactGP
        super().__init__(train_X, train_Y.squeeze(-1), GaussianLikelihood())
        self.mean_module = ConstantMean()
        self.covar_module = ScaleKernel(
            base_kernel=RBFKernel(ard_num_dims=train_X.shape[-1]),
        )
        self.to(train_X)  # make sure we're on the right device/dtype
        
    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return MultivariateNormal(mean_x, covar_x)

1.1、用Ax的BotorchModel定义一个待使用的工厂函数

Ax的BoTorchModel内部把贝叶斯优化的不同组件（模型生成、拟合、定义采集函数、优化）融入了一个功能性API中。

取决于我们想要调整哪些组件，我们可以向BoTorch模型创建器传入相关的自定义工厂函数（associated custom factory function）。为了使用一个自定义模型，我们必须实现一个模型的工厂函数，以使给定根据Ax的API说明的数据，就能实例化且拟合出一个BoTorch模型对象。

工厂函数的调用签名如下:

def get_and_fit_gpytorch_model(
    Xs: List[Tensor],
    Ys: List[Tensor],
    Yvars: List[Tensor],
    state_dict: Optional[Dict[str, Tensor]] = None,
    **kwargs: Any,
) -> Model:

此处

• Xs的第i个元素是第i个结果的训练特征，即n_i × d的张量（在我们的简单例子中只有一个结果）。
• 类似地，Ys和YVars的第i个元素是第i个结果的观测值和相关的观测方差，即n_i × 1的张量。
• state_dict是一个可选的模组状态字典，用预设的参数值初始化模型参数。

函数必须返回一个botorch模型对象。函数里如何运作由我们来定义。

使用botorch中的fit_gpytorch_model实用函数，这个简单模型的拟合就轻而易举了。（如要处理更复杂的模型，你需要使用自己的自定义拟合循环，参见Fitting a model with torch.optim）。

def _get_and_fit_simple_custom_gp(Xs, Ys, **kwargs):
    model = SimpleCustomGP(Xs[0], Ys[0])
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    fit_gpytorch_model(mll)
    return model

2、在Ax中建立优化问题

Ax的SimpleExperimentAPI需要一个可以计算实验中所有需要的指标的评价函数。这个函数需要传入一个由一组参数值构成的字典。它应该为这些指标产生一个字典，字典的键是指标名，字典的值是均值和标准误差组成的元组。

本教程中，我们使用Branin函数——一个简单的二维的人为设计的函数（译者注：关于Branin函数的详细信息参见此链接）。实际应用中，这可能是任意程度复杂的——例如，此函数可以运行一些成本高的模拟、进行A/B测试，抑或用给定参数进行机器学习模型的训练。

def branin(parameterization, *args):
    x1, x2 = parameterization["x1"], parameterization["x2"]
    y = (x2 - 5.1 / (4 * np.pi ** 2) * x1 ** 2 + 5 * x1 / np.pi - 6) ** 2
    y += 10 * (1 - 1 / (8 * np.pi)) * np.cos(x1) + 10
    # let's add some synthetic observation noise
    y += random.normalvariate(0, 0.1)
    return {"branin": (y, 0.0)}

我们需要为实验定义一个搜索空间，在搜索空间中定义参数和可行的值。

search_space = SearchSpace(
    parameters=[
        RangeParameter(
            name="x1", parameter_type=ParameterType.FLOAT, lower=-5, upper=10
        ),
        RangeParameter(
            name="x2", parameter_type=ParameterType.FLOAT, lower=-5, upper=10
        ),
    ]
)

第三步，我们创建一个SimpleExperiment，注意objective_name应该是评价函数返回的指标名称之一。

exp = SimpleExperiment(
    name="test_branin",
    search_space=search_space,
    evaluation_function=branin,
    objective_name="branin",
    minimize=True,
)

为了运行Ax优化循环中的自定义botorch模型，我们可以使用来自ax.modelbridge.factory的工厂函数get_botorch。任何传入函数的关键字参数都通往了BotorchModel创建器。为了使用我们的自定义模型，我们需要使用model_constructor参数，向get_botorch传入新定义的_get_and_fit_simple_custom_gp函数。

注意：get_botorch默认情况下会自动实施一些参数变换（如对输入normalize，对输出standardize）。这一般就是你要对连续型参数要做的标准处理工作（for standard use cases）。如果你的模型要的是原始参数（raw parameters），要确保给transform传入的是空列表transform=[]，以防任何变换。更多关于Ax中的transformations是如何运作的信息，参见Ax documentation

2.1、运行优化的循环

我们已经准备好运行贝叶斯优化的循环了。

for i in range(50):
    print(f"Running optimization batch {i + 1}/5...")
    model = get_botorch(
        experiment=exp,
        data=exp.eval(),
        search_space=exp.search_space,
        model_constructor=_get_and_fit_simple_custom_gp,
    )
    batch = exp.new_trial(generator_run=model.gen(1))
print("Done!")