项目文件夹

文件
2026-07-13 13:22:34 +08:00

38 行
1.3 KiB
Python

from pyspark.ml.connect.classification import LogisticRegression
from pyspark.ml.connect.feature import StandardScaler
from pyspark.ml.connect.pipeline import Pipeline
from pyspark.sql import SparkSession
from sklearn import datasets
import mlflow
spark = SparkSession.builder.remote("local[2]").getOrCreate()
scaler = StandardScaler(inputCol="features", outputCol="scaled_features")
lr = LogisticRegression(maxIter=10, numTrainWorkers=2, learningRate=0.001)
pipeline = Pipeline(stages=[scaler, lr])
X, y = datasets.load_iris(return_X_y=True)
spark_df = spark.createDataFrame(zip(X, y), schema="features: array<double>, label: long")
pipeline_model = pipeline.fit(spark_df)
with mlflow.start_run():
model_info = mlflow.spark.log_model(spark_model=pipeline_model, artifact_path="model")
model_uri = model_info.model_uri
print(f"Model is saved to URI: {model_uri}")
inference_df = spark_df.select("features").toPandas()
loaded_model = mlflow.spark.load_model(model_uri)
inference_result = loaded_model.transform(inference_df.copy(deep=False))
print("Loaded spark model inference result:\n")
print(inference_result)
pyfunc_model = mlflow.pyfunc.load_model(model_uri)
pyfunc_inference_result = pyfunc_model.predict(inference_df.copy(deep=False))
print("Loaded pyfunc model inference result:\n")
print(pyfunc_inference_result)