Azure ML SDK v2：如何将训练时的OHE对象用于推理部署？

核心思路是将训练时拟合完成的OHE对象保存并注册为Azure ML资产，在推理部署时加载复用，具体操作步骤如下：

1. 训练阶段保存OHE对象 #

在训练管道的脚本中，拟合OHE后，用joblib将对象保存到Azure ML指定的输出目录，该目录内容会自动上传到云端存储。

示例训练脚本（train.py）核心逻辑：

from sklearn.preprocessing import OneHotEncoder
import joblib
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--output-dir", type=str, required=True)
args = parser.parse_args()

# 假设train_data是加载后的训练数据集，categorical_cols是需要编码的列
ohe = OneHotEncoder()
ohe.fit(train_data[categorical_cols])

# 保存OHE对象到指定输出目录
joblib.dump(ohe, f"{args.output_dir}/ohe_model.joblib")

训练作业定义示例：

from azure.ai.ml import Output, command

train_job = command(
    code="./src",
    command="python train.py --output-dir ${{outputs.ohe_output}}",
    outputs={
        "ohe_output": Output(type="uri_folder", mode="rw_mount")
    },
    # 补充计算集群、环境等配置
)

2. 将OHE对象注册为Azure ML模型资产 #

训练完成后，将保存的OHE对象注册为模型资产，方便后续部署时调用。

示例代码：

from azure.ai.ml import MLClient
from azure.ai.ml.entities import Model

ml_client = MLClient.from_config()

# 从训练作业的输出路径注册OHE模型
ohe_model = Model(
    path="azureml://jobs/<你的训练作业ID>/outputs/ohe_output",
    name="onehot-encoder",
    version="1",
    description="训练阶段拟合完成的OneHotEncoder对象"
)

ml_client.models.create_or_update(ohe_model)

3. 推理评分脚本中加载OHE对象 #

编写推理评分脚本（score.py），在初始化阶段加载已注册的OHE对象，用于转换推理数据。

示例评分脚本：

import joblib
import json
import numpy as np

# 全局变量存储OHE对象
global ohe

def init():
    # Azure ML部署时，模型会被挂载到./azureml-models/<模型名称>/<版本>路径
    ohe_path = "./azureml-models/onehot-encoder/1/ohe_model.joblib"
    ohe = joblib.load(ohe_path)

def run(raw_data):
    data = json.loads(raw_data)["data"]
    # 使用OHE转换推理数据的分类特征
    transformed_features = ohe.transform(np.array(data).reshape(-1, ohe.n_features_in_))
    # 可结合预测模型输出最终结果，此处仅返回转换后的特征示例
    return {"transformed_features": transformed_features.tolist()}

4. 部署推理服务时关联OHE资产 #

部署在线推理服务时，将OHE模型资产添加到部署配置中，确保容器能挂载并加载该对象。

示例部署代码：

from azure.ai.ml.entities import Environment, ManagedOnlineDeployment, ManagedOnlineEndpoint
from azure.ai.ml.constants import AssetTypes

# 定义推理环境，需包含scikit-learn、joblib等依赖
env = Environment(
    name="ohe-inference-env",
    conda_file="./conda.yml",
    image="mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04:latest"
)

# 创建在线端点
endpoint = ManagedOnlineEndpoint(name="ohe-inference-endpoint")
ml_client.begin_create_or_update(endpoint).result()

# 创建部署，关联OHE模型资产
deployment = ManagedOnlineDeployment(
    name="ohe-deployment",
    endpoint_name="ohe-inference-endpoint",
    model=[
        Model(name="onehot-encoder", version="1", type=AssetTypes.MODEL)
        # 若有预测模型，可一并添加至此列表
    ],
    environment=env,
    instance_type="Standard_DS2_v2",
    instance_count=1
)

ml_client.begin_create_or_update(deployment).result()

额外注意事项 #

• 确保推理环境的依赖版本（如scikit-learn）与训练环境完全一致，避免加载OHE对象时出现兼容性问题。
• 若OHE与预测模型绑定使用，可将两者保存在同一目录下，注册为单个模型资产，简化部署流程。

内容的提问来源于stack exchange，提问作者amit.s