mlflow--mlflow
38 行
1.3 KiB
Python
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)
|