apache-spark - 如何在 pandas_udf 中将 Hyperopt 与 MLFlow 结合使用？

Question

我正在构建多个 Prophet 模型，其中每个模型都传递给 pandas_udf 函数，该函数训练模型并使用 MLflow 存储结果。

@pandas_udf(result_schema, PandasUDFType.GROUPED_MAP)
def forecast(data):
......
   with mlflow.start_run() as run: 
......

然后我调用这个 UDF，它为每个 KPI 训练一个模型。

df.groupBy('KPI').apply(forecast)

这个想法是，对于每个 KPI，一个模型将使用多个超参数进行训练，并将每个模型的最佳参数存储在 MLflow 中。我想使用 Hyperopt 来提高搜索效率。

在这种情况下，我应该把目标函数放在哪里？由于数据被传递到每个模型的 UDF，我想到在 UDF 中创建一个内部函数，该函数使用每次运行的数据。这有意义吗？

Answer 1

如果我没记错的话，你不能这样做，因为它类似于嵌套的 Spark 执行，而且它不适用于 Spark。你需要改变方法来处理类似的事情:

for kpi in list_of_kpis:
  run_hyperopt_tuning

如果您需要分别为每个 KPI 模型调整参数 - 因为它会分别优化参数。

如果 KPI 就像模型的超参数，那么您只需将 KPI 列表包含到搜索空间中，并在进行训练和评估的函数中加载必要的数据。