Azure ML中AutoMLStep训练模型的注册与推理问题

问题根源 #

EOFError是因为直接读取AutoMLStep输出的文件时，该文件并非可直接反序列化的模型文件，而是AutoML运行的元数据/Run引用，必须通过AutoML Run对象提取最佳模型。

步骤1：修正pipeline.py中的AutoMLStep配置 #

确保AutoMLStep输出正确的Run相关信息，为后续步骤提供获取最佳模型的入口：

from azureml.pipeline.steps import AutoMLStep
from azureml.train.automl import AutoMLConfig
from azureml.data import PipelineData
from azureml.core import Datastore

# 获取默认数据存储
default_datastore = Datastore.get(ws, 'workspaceblobstore')

# 定义AutoML训练配置
automl_config = AutoMLConfig(
    task='classification', # 根据你的任务类型调整为regression/forecasting等
    primary_metric='accuracy',
    training_data=train_dataset,
    label_column_name='label',
    n_cross_validations=5,
    iterations=10,
    model_explainability=True
)

# 定义AutoMLStep，指定输出存储AutoML运行ID
automl_step = AutoMLStep(
    name='automl_train',
    automl_config=automl_config,
    outputs=[PipelineData(name='automl_run_ref', datastore=default_datastore)],
    allow_reuse=True
)

步骤2：在register_model.py中正确获取并注册最佳模型 #

通过AutoML Run对象提取模型，避免直接读取输出文件引发EOFError：

import os
import joblib
from azureml.core import Workspace, Run, Model

# 获取当前运行上下文
run = Run.get_context()
ws = run.experiment.workspace

# 下载AutoMLStep输出的Run ID文件并读取
automl_run_file = run.input_datasets['automl_run_ref'].download()[0]
with open(automl_run_file, 'r') as f:
    automl_run_id = f.strip()

# 获取AutoML运行实例
automl_run = Run(experiment=run.experiment, run_id=automl_run_id)

# 获取最佳模型及对应运行记录
best_run, fitted_model = automl_run.get_output()

# 按非AutoML流程保存模型到本地
model_save_path = './automl_best_model.pkl'
joblib.dump(fitted_model, model_save_path)

# 注册模型到工作区
registered_model = Model.register(
    workspace=ws,
    model_path=model_save_path,
    model_name='automl-production-model',
    tags={'task': 'classification', 'training_source': 'automl_step'},
    description='AutoML训练生成的最佳分类模型'
)

# 记录模型ID到运行日志，方便后续追踪
run.log('registered_model_id', registered_model.id)

步骤3：推理阶段加载模型 #

使用Model.get_model_path加载注册后的模型，与非AutoML流程完全一致：

from azureml.core import Workspace, Model
import joblib
import pandas as pd

# 加载工作区配置
ws = Workspace.from_config()

# 获取已注册模型的本地路径
model_path = Model.get_model_path(model_name='automl-production-model')

# 加载模型
trained_model = joblib.load(model_path)

# 推理示例：对测试数据做预测
test_data = pd.read_csv('./test_data.csv')
predictions = trained_model.predict(test_data)

关键注意事项 #

• 不要直接读取AutoMLStep的输出文件内容，必须通过Run ID关联到AutoML运行实例
• automl_run.get_output()会直接返回拟合完成的模型，无需手动处理原始输出文件的反序列化
• 注册模型时确保本地保存的模型路径正确，注册后模型由工作区统一管理，可跨环境调用

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