当前位置：首页 > news >正文

Ax扩展开发指南：如何编写自定义组件和插件

news 2026/6/3 5:15:49

Ax扩展开发指南：如何编写自定义组件和插件

【免费下载链接】AxAdaptive Experimentation Platform项目地址: https://gitcode.com/gh_mirrors/ax2/Ax

Ax作为Adaptive Experimentation Platform（自适应实验平台），提供了灵活的扩展机制，让开发者能够根据特定需求定制实验组件和插件。本指南将带你快速掌握自定义Transform、Metric和Runner的核心方法，通过官方注册机制无缝集成到Ax生态系统中。

扩展开发核心概念

Ax的扩展系统基于注册机制设计，所有自定义组件需通过特定注册表注册后才能被平台识别。核心扩展点包括：

Transform：数据预处理和转换逻辑，如特征编码、归一化等
Metric：自定义评估指标计算逻辑
Runner：实验执行后端适配器
Generator：自定义采样策略实现

Ax实验组件构成示意图，展示了自定义组件如何融入整体架构

自定义Transform开发

Transform是Ax中处理实验数据的关键组件，用于实现参数编码、数据清洗和特征工程等功能。

1. 创建Transform类

首先创建继承自Transform基类的自定义转换类，实现transform_observation方法：

from ax.adapter.transforms.base import Transform from ax.core.observation import Observation class CustomLogTransform(Transform): def __init__(self, param_name: str): self.param_name = param_name def transform_observation(self, observation: Observation) -> Observation: # 对指定参数应用对数转换 modified_params = observation.parameters.copy() if self.param_name in modified_params: modified_params[self.param_name] = np.log(modified_params[self.param_name]) return observation._replace(parameters=modified_params)

2. 注册Transform

通过register_transform函数将自定义转换注册到系统：

from ax.storage.transform_registry import register_transform register_transform(CustomLogTransform)

系统注册表位于ax/storage/transform_registry.py，包含所有可用转换的元数据。

自定义Metric开发

Metric定义了实验结果的评估方式，通过实现自定义Metric可以集成特定领域的评估逻辑。

1. 实现Metric子类

创建继承自Metric的子类，实现fetch_trial_data方法：

from ax.core.metric import Metric from ax.core.types import TParameterization class LatencyMetric(Metric): def fetch_trial_data(self, trial): # 从实验结果中提取延迟数据 data = [] for arm in trial.arms: # 自定义延迟计算逻辑 latency = self._calculate_latency(arm.parameters) data.append({"arm_name": arm.name, "metric_name": self.name, "mean": latency}) return Data(data=data) def _calculate_latency(self, params: TParameterization) -> float: # 实际延迟计算实现 return ...

2. 注册Metric

使用RegistryBundle管理自定义Metric注册：

from ax.storage.registry_bundle import RegistryBundle # 创建包含自定义Metric的注册捆绑包 bundle = RegistryBundle( metric_clss={LatencyMetric: None}, # None表示自动生成类型ID runner_clss={} )

注册逻辑实现在ax/storage/registry_bundle.py，负责协调编码器、解码器和类型映射。

自定义Runner开发

Runner负责将实验配置提交到实际执行环境，如本地脚本、云服务或测试框架。

1. 实现Runner子类

from ax.core.runner import Runner from ax.core.base_trial import BaseTrial class CloudRunner(Runner): def run(self, trial: BaseTrial) -> dict[str, str]: # 将实验配置提交到云服务 job_id = self._submit_to_cloud(trial) return {"job_id": job_id} def _submit_to_cloud(self, trial: BaseTrial) -> str: # 云服务提交逻辑 return ...

2. 注册Runner

# 在RegistryBundle中注册Runner bundle = RegistryBundle( metric_clss={}, runner_clss={CloudRunner: None} )

扩展集成与使用

完成自定义组件开发后，通过AxClient加载注册捆绑包：

from ax.service.ax_client import AxClient ax_client = AxClient(registry_bundle=bundle) # 使用自定义组件定义实验 ax_client.create_experiment( name="custom_experiment", parameters=[...], objectives={"latency": Minimize()} )

集成自定义组件的Ax闭环实验流程