如果您希望充分利用 mlflow.pyfunc 的功能,并了解它可以如何在机器学习项目中被使用,这篇博文将引导您完成该过程。MLflow PyFunc 提供了创造性自由和灵活性,允许开发作为模型封装在 MLflow 中的复杂系统,这些模型遵循与传统模型相同的生命周期。本文将展示如何创建多模型设置、无缝连接数据库,以及在您的 MLflow PyFunc 模型中实现自定义的 fit 方法。
但首先,让我们用一个类比来帮助你熟悉集成模型的概念,以及为什么在下一个机器学习项目中应该考虑这一解决方案。
设想你正在市场上买房。你会仅仅根据你参观的第一套房子和某位房地产经纪人的建议就做决定吗?当然不会!买房的过程需要考虑多种因素,并从多个来源收集信息以做出明智的决定。
The house buying process explained:
- 确定您的需求: 确定您想要新房还是二手房,房屋类型、型号和建造年份。
- 研究: 寻找可用房源列表,查看折扣和优惠,阅读客户评价,并征求朋友和家人的意见。
- Evaluate: Consider the performance, location, neighborhood amenities, and price range.
简而言之,您不会直接得出结论,而是在考虑上述所有因素后再决定最佳选择。
机器学习中的集成模型基于类似的理念运行。通过将多个模型结合起来,集成学习有助于提高机器学习结果的预测性能,相较于单个模型更为优越。性能提升可能源于多种因素,例如通过对多个模型求平均来降低方差,或通过关注先前模型的错误来减少偏差。存在多种类型的集成学习技术,例如:
- 平均
- 加权平均
- 提升
然而,开发此类系统需要对集成模型的生命周期进行谨慎管理,因为整合多样的模型可能非常复杂。在这方面,MLflow PyFunc 显得非常有价值。它提供构建复杂系统的灵活性,将整个集成视为一个模型,并遵循与传统模型相同的生命周期流程。实质上,MLflow PyFunc 允许为集成模型创建定制方法,作为对诸如 scikit-learn、PyTorch 和 LangChain 等流行框架的内置 MLflow flavors 的替代方案。
项目的组成部分
注意:要重现此项目,请参阅官方 MLflow 文档,以获取有关在本地设置一个简单的 MLflow Tracking Server 的更多详细信息。
创建集成模型
与使用用于记录并与流行机器学习框架配合的内置 flavors 相比,创建 MLflow PyFunc 集成模型需要额外的步骤。
要实现一个集成模型,需要定义一个 mlflow.pyfunc 模型,这涉及创建一个继承自 PythonModel 的 Python 类并实现其构造函数和类方法。虽然基本的 PyFunc 模型只需实现 predict 方法,但集成模型需要额外的方法来管理子模型并获取多模型预测。在实例化集成模型后,必须使用自定义的 fit 方法训练集成模型的子模型。与开箱即用的 MLflow 模型类似,你需要在训练运行期间记录模型及其工件,然后在 MLflow 模型注册表中注册模型。还会向模型添加一个别名 production 以简化模型更新和推理。模型别名允许你为已注册模型的特定版本分配一个可变的命名引用。通过将别名指定到某个模型版本,就可以通过模型 URI 或模型注册表 API 轻松引用它。此设置允许在不更改服务工作负载代码的情况下无缝更新用于推理的模型版本。有关更多详细信息,请参阅 Deploy and Organize Models with Aliases and Tags。

这是 EnsembleModel 类的骨架:
import mlflow
class EnsembleModel(mlflow.pyfunc.PythonModel):
"""Ensemble model class leveraging Pyfunc for multi-model integration in MLflow."""
def __init__(self):
"""Initialize the EnsembleModel instance."""
...
def add_strategy_and_save_to_db(self):
"""Add strategies to the DuckDB database."""
...
def feature_engineering(self):
"""Perform feature engineering on input data."""
...
def initialize_models(self):
"""Initialize models and their hyperparameter grids."""
...
def fit(self):
"""Train the ensemble of models."""
...
def predict(self):
"""Predict using the ensemble of models."""
...
def load_context(self):
"""Load the preprocessor and models from the MLflow context."""
...
集成模型中的构造方法对于设置其基本要素至关重要。它建立了诸如 preprocessor、用于存储训练模型的字典、指向 DuckDB 数据库的路径,以及用于管理不同集成策略的 pandas DataFrame 等关键属性。此外,它利用 initialize_models 方法来定义集成中包含的子模型。
import pandas as pd
def __init__(self):
"""
Initializes the EnsembleModel instance.
Sets up an empty preprocessing pipeline, a dictionary for fitted models,
and a DataFrame to store strategies. Also calls the method to initialize sub-models.
"""
self.preprocessor = None
self.fitted_models = {}
self.db_path = None
self.strategies = pd.DataFrame(columns=["strategy", "model_list", "weights"])
self.initialize_models()
自定义的 add_strategy_and_save_to_db 方法能够向模型添加新的集成策略,并将其存储到 DuckDB 数据库中。该方法接受一个包含策略的 pandas DataFrame 和数据库路径作为输入。它将新的策略追加到已有策略中,并保存到在集成模型初始化时指定的数据库中。该方法便于管理多种集成策略,并确保它们为将来使用而被持久化存储。
import duckdb
import pandas as pd
def add_strategy_and_save_to_db(self, strategy_df: pd.DataFrame, db_path: str) -> None:
"""Add strategies from a DataFrame and save them to the DuckDB database.
Args:
strategy_df (pd.DataFrame): DataFrame containing strategies.
db_path (str): Path to the DuckDB database.
"""
# Update the instance-level database path for the current object
self.db_path = db_path
# Attempt to concatenate new strategies with the existing DataFrame
try:
self.strategies = pd.concat([self.strategies, strategy_df], ignore_index=True)
except Exception as e:
# Print an error message if any exceptions occur during concatenation
print(f"Error concatenating DataFrames: {e}")
return # Exit early to prevent further errors
# Use context manager for the database connection
try:
with duckdb.connect(self.db_path) as con:
# Register the strategies DataFrame as a temporary table in DuckDB
con.register("strategy_df", self.strategies)
# Drop any existing strategies table and create a new one with updated strategies
con.execute("DROP TABLE IF EXISTS strategies")
con.execute("CREATE TABLE strategies AS SELECT * FROM strategy_df")
except Exception as e:
# Print an error message if any exceptions occur during database operations
print(f"Error executing database operations: {e}")
下面的示例演示了如何使用此方法将策略添加到数据库。
import pandas as pd
# Initialize ensemble model
ensemble_model = EnsembleModel()
# Define strategies for the ensemble model
strategy_data = {
"strategy": ["average_1"],
"model_list": ["random_forest,xgboost,decision_tree,gradient_boosting,adaboost"],
"weights": ["1"],
}
# Create a DataFrame to hold the strategy information
strategies_df = pd.DataFrame(strategy_data)
# Add strategies to the database
ensemble_model.add_strategy_and_save_to_db(strategies_df, "models/strategies.db")
该 DataFrame strategy_data 包含:
- strategy: 模型预测所用策略的名称。
- weights: 一个以逗号分隔的权重列表,分配给
model_list中的每个模型。如果未提供,则表示相等权重或默认值。
| average_1 | random_forest,xgboost,decision_tree,gradient_boosting,adaboost | 1 |
The feature_engineering 方法通过处理缺失值、对数值特征进行缩放以及对类别特征进行编码来预处理输入数据。它对数值和类别特征分别应用不同的变换,并以 NumPy 数组的形式返回处理后的特征。该方法对于将数据准备为适合模型训练的格式至关重要,能够确保一致性并提升模型性能。
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
def feature_engineering(self, X: pd.DataFrame) -> np.ndarray:
"""
Applies feature engineering to the input data X, including imputation, scaling, and encoding.
Args:
X (pd.DataFrame): Input features with potential categorical and numerical columns.
Returns:
np.ndarray: Processed feature array after transformations.
"""
# Convert columns with 'object' dtype to 'category' dtype for proper handling of categorical features
X = X.apply(
lambda col: col.astype("category") if col.dtypes == "object" else col
)
# Identify categorical and numerical features from the DataFrame
categorical_features = X.select_dtypes(include=["category"]).columns
numerical_features = X.select_dtypes(include=["number"]).columns
# Define the pipeline for numerical features: imputation followed by scaling
numeric_transformer = Pipeline(
steps=[
(
"imputer",
SimpleImputer(strategy="median"),
), # Replace missing values with the median
(
"scaler",
StandardScaler(),
), # Standardize features by removing the mean and scaling to unit variance
]
)
# Define the pipeline for categorical features: imputation followed by one-hot encoding
categorical_transformer = Pipeline(
steps=[
(
"imputer",
SimpleImputer(strategy="most_frequent"),
), # Replace missing values with the most frequent value
(
"onehot",
OneHotEncoder(handle_unknown="ignore"),
), # Encode categorical features as a one-hot numeric array
]
)
# Create a ColumnTransformer to apply the appropriate pipelines to the respective feature types
preprocessor = ColumnTransformer(
transformers=[
(
"num",
numeric_transformer,
numerical_features,
), # Apply the numeric pipeline to numerical features
(
"cat",
categorical_transformer,
categorical_features,
), # Apply the categorical pipeline to categorical features
]
)
# Fit and transform the input data using the preprocessor
X_processed = preprocessor.fit_transform(X)
# Store the preprocessor for future use in the predict method
self.preprocessor = preprocessor
return X_processed
初始化模型
initialize_models 方法设置了一个包含各种机器学习模型及其超参数网格的字典。 这包括诸如 RandomForest、XGBoost、DecisionTree、GradientBoosting 和 AdaBoost 等模型。 这一步对于准备集成的子模型并指定在训练期间需要调整的超参数至关重要,从而确保每个模型被正确配置并准备好进行训练。
from sklearn.ensemble import (
AdaBoostRegressor,
GradientBoostingRegressor,
RandomForestRegressor,
)
from sklearn.tree import DecisionTreeRegressor
from xgboost import XGBRegressor
def initialize_models(self) -> None:
"""
Initializes a dictionary of models along with their hyperparameter grids for grid search.
"""
# Define various regression models with their respective hyperparameter grids for tuning
self.models = {
"random_forest": (
RandomForestRegressor(random_state=42),
{"n_estimators": [50, 100, 200], "max_depth": [None, 10, 20]},
),
"xgboost": (
XGBRegressor(random_state=42),
{"n_estimators": [50, 100, 200], "max_depth": [3, 6, 10]},
),
"decision_tree": (
DecisionTreeRegressor(random_state=42),
{"max_depth": [None, 10, 20]},
),
"gradient_boosting": (
GradientBoostingRegressor(random_state=42),
{"n_estimators": [50, 100, 200], "max_depth": [3, 5, 7]},
),
"adaboost": (
AdaBoostRegressor(random_state=42),
{"n_estimators": [50, 100, 200], "learning_rate": [0.01, 0.1, 1.0]},
),
}
定义自定义 fit 方法来训练并保存多模型
import os
import joblib
import pandas as pd
from sklearn.model_selection import GridSearchCV
def fit(
self, X_train_processed: pd.DataFrame, y_train: pd.Series, save_path: str
) -> None:
"""
Trains the ensemble of models using the provided preprocessed training data.
Args:
X_train_processed (pd.DataFrame): Preprocessed feature matrix for training.
y_train (pd.Series): Target variable for training.
save_path (str): Directory path where trained models will be saved.
"""
# Create the directory for saving models if it does not exist
os.makedirs(save_path, exist_ok=True)
# Iterate over each model and its parameter grid
for model_name, (model, param_grid) in self.models.items():
# Perform GridSearchCV to find the best hyperparameters for the current model
grid_search = GridSearchCV(
model, param_grid, cv=5, n_jobs=-1, scoring="neg_mean_squared_error"
)
grid_search.fit(
X_train_processed, y_train
) # Fit the model with the training data
# Save the best estimator from GridSearchCV
best_model = grid_search.best_estimator_
self.fitted_models[model_name] = best_model
# Save the trained model to disk
joblib.dump(best_model, os.path.join(save_path, f"{model_name}.pkl"))
为简化推理流程,每个 PyFunc 模型都应定义一个自定义的 predict 方法,作为推理的唯一入口点。这种做法在推理时将模型的内部工作抽象出来,无论是自定义的 PyFunc 模型,还是面向流行机器学习框架的开箱即用的 MLflow 内置 flavor,都适用。
集成模型的自定义 predict 方法旨在收集并合并子模型的预测,支持各种聚合策略(例如,平均、加权)。该过程包括以下步骤:
- 加载基于用户自定义方法的子模型预测聚合策略。
- 预处理输入数据。
- 根据指定的策略汇总模型预测。
import duckdb
import joblib
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
def predict(self, context, model_input: pd.DataFrame) -> np.ndarray:
"""
Predicts the target variable using the ensemble of models based on the selected strategy.
Args:
context: MLflow context object.
model_input (pd.DataFrame): Input features for prediction.
Returns:
np.ndarray: Array of predicted values.
Raises:
ValueError: If the strategy is unknown or no models are fitted.
"""
# Check if the 'strategy' column is present in the input DataFrame
if "strategy" in model_input.columns:
# Extract the strategy and drop it from the input features
print(f"Strategy: {model_input['strategy'].iloc[0]}")
strategy = model_input["strategy"].iloc[0]
model_input.drop(columns=["strategy"], inplace=True)
else:
# Default to 'average' strategy if none is provided
strategy = "average"
# Load the strategy details from the pre-loaded strategies DataFrame
loaded_strategy = self.strategies[self.strategies["strategy"] == strategy]
if loaded_strategy.empty:
# Raise an error if the specified strategy is not found
raise ValueError(
f"Strategy '{strategy}' not found in the pre-loaded strategies."
)
# Parse the list of models to be used for prediction
model_list = loaded_strategy["model_list"].iloc[0].split(",")
# Transform input features using the preprocessor, if available
if self.preprocessor is None:
# Feature engineering is required if the preprocessor is not set
X_processed = self.feature_engineering(model_input)
else:
# Use the existing preprocessor to transform the features
X_processed = self.preprocessor.transform(model_input)
if not self.fitted_models:
# Raise an error if no models are fitted
raise ValueError("No fitted models found. Please fit the models first.")
# Collect predictions from all models specified in the strategy
predictions = np.array(
[self.fitted_models[model].predict(X_processed) for model in model_list]
)
# Apply the specified strategy to combine the model predictions
if "average" in strategy:
# Calculate the average of predictions from all models
return np.mean(predictions, axis=0)
elif "weighted" in strategy:
# Extract weights from the strategy and normalize them
weights = [float(w) for w in loaded_strategy["weights"].iloc[0].split(",")]
weights = np.array(weights)
weights /= np.sum(weights) # Ensure weights sum to 1
# Compute the weighted average of predictions
return np.average(predictions, axis=0, weights=weights)
else:
# Raise an error if an unknown strategy is encountered
raise ValueError(f"Unknown strategy: {strategy}")
这个初始化过程包括:
- 使用包含工件引用的 context object 加载模型工件,包括预训练模型和预处理器。
- 正在从 DuckDB 数据库获取策略定义。
import duckdb
import joblib
import pandas as pd
def load_context(self, context) -> None:
"""
Loads the preprocessor and models from the MLflow context.
Args:
context: MLflow context object which provides access to saved artifacts.
"""
# Load the preprocessor if its path is specified in the context artifacts
preprocessor_path = context.artifacts.get("preprocessor", None)
if preprocessor_path:
self.preprocessor = joblib.load(preprocessor_path)
# Load each model from the context artifacts and store it in the fitted_models dictionary
for model_name in self.models.keys():
model_path = context.artifacts.get(model_name, None)
if model_path:
self.fitted_models[model_name] = joblib.load(model_path)
else:
# Print a warning if a model is not found in the context artifacts
print(
f"Warning: {model_name} model not found in artifacts. Initialized but not fitted."
)
# Reconnect to the DuckDB database to load the strategies
conn = duckdb.connect(self.db_path)
# Fetch strategies from the DuckDB database into the strategies DataFrame
self.strategies = conn.execute("SELECT * FROM strategies").fetchdf()
# Close the database connection
conn.close()
将这一切结合在一起
您可以使用Creating the Ensemble Model部分提供的骨架来组装整个EnsembleModel类。每个方法都演示了其包含的具体依赖项。现在,您只需按照给定的大纲将这些方法合并到类定义中。可以随意添加任何适合您特定用例或增强集成模型功能的自定义逻辑。
在将所有内容封装到 PyFunc 模型之后,集成模型的生命周期与传统的 MLflow 模型非常相似。下图描述了该模型的生命周期。

MLflow 跟踪
在对数据进行预处理之后,我们使用自定义的 fit 方法来训练集成模型中的所有子模型。该方法对每个子模型应用网格搜索以找到最佳超参数,拟合到训练数据,并保存训练好的模型以备将来使用。
import datetime
import os
import joblib
import mlflow
import pandas as pd
from mlflow.models.signature import infer_signature
from sklearn.model_selection import train_test_split
# Initialize the MLflow client
client = mlflow.MlflowClient()
# Set the URI of your MLflow Tracking Server
remote_server_uri = "..." # Replace with your server URI
# Point MLflow to your MLflow Tracking Server
mlflow.set_tracking_uri(remote_server_uri)
# Set the experiment name for organizing runs in MLflow
mlflow.set_experiment("Ensemble Model")
# Load dataset from the provided URL
data = pd.read_csv(
"https://github.com/zobi123/Machine-Learning-project-with-Python/blob/master/Housing.csv?raw=true"
)
# Separate features and target variable
X = data.drop("price", axis=1)
y = data["price"]
# Split dataset into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)
# Create a directory to save the models and related files
os.makedirs("models", exist_ok=True)
# Initialize and train the EnsembleModel
ensemble_model = EnsembleModel()
# Preprocess the training data using the defined feature engineering method
X_train_processed = ensemble_model.feature_engineering(X_train)
# Fit the models with the preprocessed training data and save them
ensemble_model.fit(X_train_processed, y_train, save_path="models")
# Infer the model signature using a small example from the training data
example_input = X_train[:1] # Use a single sample for signature inference
example_input["strategy"] = "average"
example_output = y_train[:1]
signature = infer_signature(example_input, example_output)
# Save the preprocessing pipeline to disk
joblib.dump(ensemble_model.preprocessor, "models/preprocessor.pkl")
# Define strategies for the ensemble model
strategy_data = {
"strategy": [
"average_1",
"average_2",
"weighted_1",
"weighted_2",
"weighted_3",
"weighted_4",
],
"model_list": [
"random_forest,xgboost,decision_tree,gradient_boosting,adaboost",
"decision_tree",
"random_forest,xgboost,decision_tree,gradient_boosting,adaboost",
"random_forest,xgboost,gradient_boosting",
"decision_tree,adaboost",
"xgboost,gradient_boosting",
],
"weights": ["1", "1", "0.2,0.3,0.1,0.2,0.2", "0.4,0.4,0.2", "0.5,0.5", "0.7,0.3"],
}
# Create a DataFrame to hold the strategy information
strategies_df = pd.DataFrame(strategy_data)
# Add strategies to the database
ensemble_model.add_strategy_and_save_to_db(strategies_df, "models/strategies.db")
# Define the Conda environment configuration for the MLflow model
conda_env = {
"name": "mlflow-env",
"channels": ["conda-forge"],
"dependencies": [
"python=3.8",
"scikit-learn=1.3.0",
"xgboost=2.0.3",
"joblib=1.2.0",
"pandas=1.5.3",
"numpy=1.23.5",
"duckdb=1.0.0",
{
"pip": [
"mlflow==2.14.1",
]
},
],
}
# Get current timestamp
timestamp = datetime.datetime.now().isoformat()
# Log the model using MLflow
with mlflow.start_run(run_name=timestamp) as run:
# Log parameters, artifacts, and model signature
mlflow.log_param("model_type", "EnsembleModel")
artifacts = {
model_name: os.path.join("models", f"{model_name}.pkl")
for model_name in ensemble_model.models.keys()
}
artifacts["preprocessor"] = os.path.join("models", "preprocessor.pkl")
artifacts["strategies_db"] = os.path.join("models", "strategies.db")
mlflow.pyfunc.log_model(
artifact_path="ensemble_model",
python_model=ensemble_model,
artifacts=artifacts,
conda_env=conda_env,
signature=signature,
)
print(f"Model logged in run {run.info.run_id}")
在 MLflow 中注册模型
在模型训练完成后,下一步是将集成模型注册到 MLflow。此过程涉及将训练好的模型、预处理管道和相关策略记录到 MLflow Tracking Server。这确保了集成模型的所有组件被系统地保存并进行版本控制,从而促进可复现性和可追溯性。
此外,我们会为该模型的初始版本指定一个生产别名。此指定确立了一个基线模型,可用于评估未来的迭代。通过将该版本标记为 production 模型,我们能够有效地对改进进行基准测试,并确认后续版本在该既定基线上确实取得了可衡量的进步。
# Register the model in MLflow and assign a production alias
model_uri = f"runs:/{run.info.run_id}/ensemble_model"
model_details = mlflow.register_model(model_uri=model_uri, name="ensemble_model")
client.set_registered_model_alias(
name="ensemble_model", alias="production", version=model_details.version
)
下图展示了在 MLflow UI 中我们集成模型到目前为止的完整生命周期。

使用 predict 方法执行推理
import pandas as pd
import mlflow
from sklearn.metrics import r2_score
# Load the registered model using its alias
loaded_model = mlflow.pyfunc.load_model(
model_uri=f"models:/ensemble_model@production"
)
# Define the different strategies for evaluation
strategies = [
"average_1",
"average_2",
"weighted_1",
"weighted_2",
"weighted_3",
"weighted_4",
]
# Initialize a DataFrame to store the results of predictions
results_df = pd.DataFrame()
# Iterate over each strategy, make predictions, and calculate R^2 scores
for strategy in strategies:
# Create a test DataFrame with the current strategy
X_test_with_params = X_test.copy()
X_test_with_params["strategy"] = strategy
# Use the loaded model to make predictions
y_pred = loaded_model.predict(X_test_with_params)
# Calculate R^2 score for the predictions
r2 = r2_score(y_test, y_pred)
# Store the results and R^2 score in the results DataFrame
results_df[strategy] = y_pred
results_df[f"r2_{strategy}"] = r2
# Add the actual target values to the results DataFrame
results_df["y_test"] = y_test.values
和开箱即用的 MLflow 模型类似,您首先使用 mlflow.pyfunc.load_model 加载集成模型来生成房价预测。定义用于汇总子模型预测的不同策略并创建同时包含房屋数据特征和汇总策略的模型输入后,只需调用集成模型的 predict 方法即可获得汇总后的房价预测。
使用不同策略评估模型性能
为了评估我们集成模型的性能,我们计算了不同聚合策略下的平均R²分数。这些策略包括对子模型的简单平均以及加权组合,模型及其各自的权重配置各不相同。通过比较R²分数,我们可以评估哪些策略提供最准确的预测。
下面的条形图显示了每种策略的平均 R² 分数。较高的数值表示更好的预测性能。如图所示,集成策略通常优于单一模型,正如我们的第二种依赖单个 DecisionTree (average_2) 的策略所示,这证明了聚合多个子模型预测的有效性。此可视化比较突显了使用集成方法的好处,尤其是通过加权策略来优化每个子模型贡献时。

不仅提供实现预期结果的结构化方法,本博文还基于实践经验展示了可行的方法,并讨论了潜在的挑战及其解决办法。
附加资源
探索以下资源以更深入地了解 MLflow PyFunc 模型:



