使用 AutoML Python API 训练分类模型

此示例笔记本演示如何使用 AutoML Python API 在 Databricks 上训练分类模型。 使用 UCI 人口普查收入数据集,调用 automl.classify() 预测个人年收入是否超过 5万美元,然后使用最佳实验对 pandas 和 Spark 数据帧进行推理。

要求

用于机器学习的 Databricks Runtime。

人口普查收入数据集

此数据集包含 1994 年人口普查数据库中的人口普查数据。 每行表示一组个人。 目标是确定一个群体的年收入是否超过 5 万。 此分类在>50k列中表示为字符串,其值为<=50K

from pyspark.sql.types import DoubleType, StringType, StructType, StructField

schema = StructType([
  StructField("age", DoubleType(), False),
  StructField("workclass", StringType(), False),
  StructField("fnlwgt", DoubleType(), False),
  StructField("education", StringType(), False),
  StructField("education_num", DoubleType(), False),
  StructField("marital_status", StringType(), False),
  StructField("occupation", StringType(), False),
  StructField("relationship", StringType(), False),
  StructField("race", StringType(), False),
  StructField("sex", StringType(), False),
  StructField("capital_gain", DoubleType(), False),
  StructField("capital_loss", DoubleType(), False),
  StructField("hours_per_week", DoubleType(), False),
  StructField("native_country", StringType(), False),
  StructField("income", StringType(), False)
])
input_df = spark.read.format("csv").schema(schema).load("/databricks-datasets/adult/adult.data")

训练/测试拆分

train_df, test_df = input_df.randomSplit([0.99, 0.01], seed=42)
display(train_df)

Training

以下命令启动 AutoML 运行。 您必须在 target_col 参数中提供模型应预测的列。
运行完成后,可以按照指向最佳试用笔记本的链接来检查训练代码。 此笔记本还包括特征重要性图。

from databricks import automl
summary = automl.classify(train_df, target_col="income", timeout_minutes=30)

以下命令显示有关 AutoML 输出的信息。

help(summary)

推断

可以使用 AutoML 训练的模型对新数据进行预测。 以下示例演示如何对 pandas DataFrame 中的数据进行预测,或将模型注册为 Spark UDF,以便对 Spark 数据帧进行预测。

model_uri = summary.best_trial.model_path
# model_uri = "<model-uri-from-generated-notebook>"

pandas 数据框架 (DataFrame)

import mlflow

# Prepare test dataset
test_pdf = test_df.toPandas()
y_test = test_pdf["income"]
X_test = test_pdf.drop("income", axis=1)

# Run inference using the best model
model = mlflow.pyfunc.load_model(model_uri)
predictions = model.predict(X_test)
test_pdf["income_predicted"] = predictions
display(test_pdf)

Spark 数据帧

predict_udf = mlflow.pyfunc.spark_udf(spark, model_uri=model_uri, result_type="string")
display(test_df.withColumn("income_predicted", predict_udf()))

测试

使用最终模型对留出测试集进行预测,以估算模型在生产环境中的表现如何。 此图显示了正确预测和错误预测之间的细分。

import sklearn.metrics

model = mlflow.sklearn.load_model(model_uri)
sklearn.metrics.plot_confusion_matrix(model, X_test, y_test)

注册并部署模型

可以像 MLflow 模型注册表中的其他任何模型一样注册和部署 AutoML 训练的模型。 请参阅 日志、加载和注册 MLflow 模型

故障 排除: No module named pandas.core.indexes.numeric

使用 Mosaic AI 模型服务来提供 AutoML 训练的模型时,您可能会看到错误 No module named pandas.core.indexes.numeric。 当 AutoML 使用的版本与模型服务终结点环境中的版本不同时,会发生这种情况 pandas 。 若要解决问题,请执行以下操作:

  1. 下载 add-pandas-dependency.py 脚本。 脚本编辑 requirements.txtconda.yaml 记录的模型要固定 pandas==1.5.3
  2. 编辑脚本以纳入模型记录所在的 MLflow 运行 run_id
  3. 重新注册模型。
  4. 部署新的模型版本。

示例笔记本

使用 AutoML Python API 训练分类模型

获取笔记本

后续步骤

AutoML Python API 参考