如何使用AX运行并发试用？

Question

根据AX文档（https://ax.dev/docs/api.html），在Service API下，它表示“可以并行评估试用版”。我该如何实现？

我有4个GPU，我想一次评估4组不同的超参数（同时进行4次训练）。

from ax.service.ax_client import AxClient
from ax.utils.measurement.synthetic_functions import branin

ax = AxClient()
ax.create_experiment(
    name="branin_test_experiment",
    parameters=[
        {
            "name": "x1",
            "type": "range",
            "bounds": [-5.0, 10.0],
            "value_type": "float",
        },
        {
            "name": "x2",
            "type": "range",
            "bounds": [0.0, 10.0],
        },
    ],
    objective_name="branin",
    minimize=True,
)

for _ in range(15):
    parameters, trial_index = ax.get_next_trial()
    ax.complete_trial(trial_index=trial_index, raw_data=branin(parameters["x1"], parameters["x2"]))

best_parameters, metrics = ax.get_best_parameters()