如何向MLflow部署的CNN模型服务端点发送原始输入图像？

1. 构建自定义PyFunc包装类 #

MLflow的mlflow.pyfunc.PythonModel支持自定义预测逻辑，你可以把预处理逻辑和原始分类模型都封装在这个类中，实现端到端的预测流程。示例代码如下：

import mlflow.pyfunc
import numpy as np
from PIL import Image
import io

class CatDogClassifierWrapper(mlflow.pyfunc.PythonModel):
    # 初始化时传入预处理所需参数，如目标尺寸、归一化系数
    def __init__(self, target_size=(224,224), normalize_scale=1/255.0):
        self.target_size = target_size
        self.normalize_scale = normalize_scale
        self.raw_model = None  # 存储实际执行预测的猫狗分类模型实例
    
    # 模型加载时自动执行，从传入的artifact中读取你训练好的原始模型
    def load_context(self, context):
        # 以下为TensorFlow/Keras模型加载示例，PyTorch模型替换为对应加载逻辑即可
        import tensorflow as tf
        self.raw_model = tf.keras.models.load_model(context.artifacts["base_model"])
    
    # 对外暴露的预测接口，所有服务端点收到的请求都会进入该方法
    def predict(self, context, model_input):
        processed_imgs = []
        # 遍历请求传入的原始图片数据，执行预处理，逻辑和你本地send_request.py中的预处理完全一致
        for img_bytes in model_input.iloc[:, 0]:
            img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
            img = img.resize(self.target_size)
            img_arr = np.array(img) * self.normalize_scale
            processed_imgs.append(img_arr)
        
        # 组装为模型要求的批次输入格式
        input_batch = np.stack(processed_imgs, axis=0)
        # 执行预测
        predictions = self.raw_model.predict(input_batch)
        # 可根据需求自定义返回格式，比如直接返回分类标签和置信度，无需客户端再做解析
        return [{"label": "cat" if pred[0]<0.5 else "dog", "confidence": float(pred[0])} for pred in predictions]

2. 记录自定义模型到MLflow #

训练完成后，将原始模型作为附件和自定义包装类一起记录到MLflow，同时配置模型签名和运行依赖：

from mlflow.models.signature import ModelSignature
from mlflow.types.schema import Schema, ColSpec, BinaryType, StringType, FloatType

# 定义模型签名：输入为二进制格式的原始图片，输出为分类标签和置信度
input_schema = Schema([ColSpec(BinaryType(), "image")])
output_schema = Schema([
    ColSpec(StringType(), "label"),
    ColSpec(FloatType(), "confidence")
])
signature = ModelSignature(inputs=input_schema, outputs=output_schema)

# 声明附件，填写你本地已经训练好的原始模型存储路径
artifacts = {
    "base_model": "./本地原始猫狗分类模型的存储路径"
}

# 声明运行环境依赖，包含所有需要用到的第三方包
conda_env = {
    "channels": ["defaults"],
    "dependencies": [
        "python=3.9",
        "pip",
        {
            "pip": [
                "mlflow>=2.0",
                "tensorflow>=2.10", # 按需替换为pytorch等对应框架版本
                "pillow>=9.0",
                "numpy>=1.21"
            ]
        }
    ]
}

# 记录模型到MLflow
with mlflow.start_run() as run:
    mlflow.pyfunc.log_model(
        artifact_path="cat_dog_classifier",
        python_model=CatDogClassifierWrapper(target_size=(224,224)),
        artifacts=artifacts,
        conda_env=conda_env,
        signature=signature
    )
    print(f"模型记录完成，run_id: {run.info.run_id}")

3. 部署服务并测试 #

使用标准mlflow serve命令部署刚才记录的自定义模型即可，请求时直接传入原始图片，无需提前做预处理，curl测试示例如下：

curl -X POST -H "Content-Type:multipart/form-data" -F "image=@./test_cat.jpg" http://127.0.0.1:5000/invocations

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