如何在 mlflow 中添加系数、p 值和相关变量名称？

Question

我正在运行一个线性回归模型，我想将每个变量的系数和 P 值以及变量名称添加到 mlflow 输出的指标中。我是使用 mlflow 的新手，并且对此不太熟悉。下面是我的部分代码示例

with mlflow.start_run(run_name=p_key + '_' + str(o_key)):
    
    lr = LinearRegression(
      featuresCol = 'features',
      labelCol = target_var,
      maxIter = 10,
      regParam = 0.0,
      elasticNetParam = 0.0,
      solver="normal"
        )
    
    lr_model_item = lr.fit(train_model_data)
    lr_coefficients_item = lr_model_item.coefficients
    lr_coefficients_intercept = lr_model_item.intercept
    
    lr_predictions_item = lr_model_item.transform(train_model_data)
    lr_predictions_item_oos = lr_model_item.transform(test_model_data)
    
    rsquared = lr_model_item.summary.r2
    
    # Log mlflow attributes for mlflow UI
    mlflow.log_metric("rsquared", rsquared)
    mlflow.log_metric("intercept", lr_coefficients_intercept)
    for i in lr_coefficients_item:
      mlflow.log_metric('coefficients', lr_coefficients_item[i])

想知道这是否可行？在最终输出中，我应该有截距、系数、p 值和相关的变量名称。

Answer 1

如果我理解正确，您想在 MLFlow 中分别注册每个变量名称的 p 值和系数。 Spark ML 的难点在于，在将所有列传递给给定的估计器（例如LinearRegression）之前，通常会将所有列合并到一个“特征”列中。因此，人们可以忽略哪个名称属于哪个列。

我们可以通过定义以下函数 [1] 从您的线性模型中获取“特征”列中每个特征的名称：

from itertools import chain

def feature_names(model, df):
  features_dict = df.schema[model.summary.featuresCol].metadata["ml_attr"]["attrs"].values()
  return sorted([(attr["idx"], attr["name"]) for attr in chain(*features_dict)])

上面的函数返回一个包含元组列表的排序列表，其中第一个条目对应于“features”列中特征的索引，第二个条目对应实际特征的名称。

通过在您的代码中使用上述函数，我们现在可以轻松地将特征名称与“特征”列中的列匹配，从而注册每个特征的系数和 p 值。

def has_pvalue(model):
  ''' Check if the given model supports pValues associated '''
  try:
    model.summary.pValues
    return True
  except:
    return False


with mlflow.start_run():
  lr = LinearRegression(
    featuresCol="features",
    labelCol="label",
    maxIter = 10,
    regParam = 1.0,
    elasticNetParam = 0.0,
    solver = "normal"
  )
  lr_model = lr.fit(train_data)

  mlflow.log_metric("rsquared", lr_model.summary.r2)
  mlflow.log_metric("intercept", lr_model.intercept)
  
  for index, name in feature_names(lr_model, train_data):
    mlflow.log_metric(f"Coef. {name}", lr_model.coefficients[index])
    if has_pvalue(lr_model):
      # P-values are not always available. This depends on the model configuration.
      mlflow.log_metric(f"P-val. {name}", lr_model.summary.pValues[index])

[1]：Related Stackoverflow question