MLFlow：如何从已加载模型中获取额外方法？

使用场景 #

定义了mlflow.pyfunc.PythonModel的子类，添加了静态方法和实例方法用于将预测结果解析为不同格式。从MLflow模型仓库加载模型后，希望调用这些自定义辅助方法，而不仅仅是通用的predict方法。

示例代码 #

自定义模型类 #

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        prediction = # 执行预测逻辑
        return prediction

    @staticmethod
    def parse_prediction_to_format_x(prediction):
        prediction_formatted = # 格式X解析逻辑
        return prediction_formatted

    def parse_prediction_to_format_y(self, prediction):
        prediction_formatted = # 格式Y解析逻辑
        return prediction_formatted

加载模型代码 #

loaded_model = mlflow.pyfunc.load_model(
    model_uri=saved_model_path.absolute().as_uri()
)

# 通用predict方法可正常调用
predicted = loaded_model.predict(input_data)

核心问题 #

加载后的loaded_model无法直接调用自定义的parse_prediction_to_format_x（静态方法）和parse_prediction_to_format_y（实例方法），该如何访问这些方法？

解决方案 #

方法1：访问MLflow内部存储的原始模型实例 #

MLflow加载后的PyFuncModel对象会将原始模型实例存储在_model_impl属性中，可直接通过该属性调用自定义方法：

predicted = loaded_model.predict(input_data)

# 调用实例方法
formatted_y = loaded_model._model_impl.parse_prediction_to_format_y(predicted)

# 调用静态方法（通过实例调用即可）
formatted_x = loaded_model._model_impl.parse_prediction_to_format_x(predicted)

注意：_model_impl是MLflow内部属性，后续版本可能发生变更，使用时需留意兼容性。

方法2：扩展predict方法支持格式化参数 #

在模型类中新增带格式参数的predict方法，将预测和格式化逻辑整合：

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, format_type=None):
        prediction = # 执行预测逻辑
        if format_type == "x":
            return self.parse_prediction_to_format_x(prediction)
        elif format_type == "y":
            return self.parse_prediction_to_format_y(prediction)
        else:
            return prediction

    @staticmethod
    def parse_prediction_to_format_x(prediction):
        prediction_formatted = # 格式X解析逻辑
        return prediction_formatted

    def parse_prediction_to_format_y(self, prediction):
        prediction_formatted = # 格式Y解析逻辑
        return prediction_formatted

加载后直接调用：

# 获取格式X的结果
formatted_x = loaded_model.predict(input_data, format_type="x")

# 获取格式Y的结果
formatted_y = loaded_model.predict(input_data, format_type="y")

# 获取原始预测结果
predicted = loaded_model.predict(input_data)

方法3：将格式化逻辑独立为工具模块并打包 #

把格式化方法抽离到单独的工具模块，保存模型时确保该模块被包含在模型依赖中：

1. 创建工具模块prediction_formatters.py：

def parse_prediction_to_format_x(prediction):
    prediction_formatted = # 格式X解析逻辑
    return prediction_formatted

def parse_prediction_to_format_y(prediction):
    prediction_formatted = # 格式Y解析逻辑
    return prediction_formatted

2. 模型类中引用该模块：

from prediction_formatters import parse_prediction_to_format_x, parse_prediction_to_format_y

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        prediction = # 执行预测逻辑
        return prediction

3. 保存模型时确保prediction_formatters.py被包含（可通过MLflow的code_path参数指定），加载模型后直接导入使用：

from prediction_formatters import parse_prediction_to_format_x, parse_prediction_to_format_y

predicted = loaded_model.predict(input_data)
formatted_x = parse_prediction_to_format_x(predicted)
formatted_y = parse_prediction_to_format_y(predicted)

内容的提问来源于stack exchange，提问作者YFl