能否无需Pickle序列化模型即可在MLflow模型注册表中注册？

MLflow注册模型必须指向符合MLflow模型格式的目录，这个目录的核心标志是包含MLmodel元数据文件——它定义了模型的flavor、依赖、artifact路径等关键信息。你之前直接注册单个model_state.pth文件，不符合MLflow的模型格式要求，因此注册表不会生成对应的模型版本。

你不需要序列化整个模型，只需要手动构造一个符合MLflow格式的模型目录（包含MLmodel文件和你的state_dict），或者借助MLflow的pyfunc flavor自动生成合法格式，就能完成注册。下面是两种可行方案：

方案1：手动构造MLflow模型目录 #

1. 创建一个模型目录（比如custom_model），将你的model_state.pth放入其中
2. 在目录内创建MLmodel文件，写入基础元数据（因为你不需要load_model，可以简化内容）
3. 上传整个目录作为artifact，再注册这个目录的URI

具体代码 #

import os
import torch
import mlflow

# 1. 创建模型目录并保存state_dict
model_dir = "custom_model"
os.makedirs(model_dir, exist_ok=True)
state_dict_path = os.path.join(model_dir, "model_state.pth")
torch.save(cli.model.state_dict(), state_dict_path)

# 2. 手动写入MLmodel文件
mlmodel_content = """
artifact_path: custom_model
flavors:
  pytorch:
    model_data: model_state.pth
    pytorch_version: "2.0.1"  # 替换为你实际使用的PyTorch版本
"""
with open(os.path.join(model_dir, "MLmodel"), "w") as f:
    f.write(mlmodel_content.strip())

# 3. 上传目录并注册模型
with mlflow.start_run() as run:
    run_id = run.info.run_id
    mlflow.log_artifacts(model_dir, artifact_path="custom_model")
    model_uri = f"runs:/{run_id}/custom_model"
    mlflow.register_model(model_uri, "Test")

方案2：借助MLflow PyFunc自动生成格式 #

用MLflow的pyfunc flavor可以自动生成MLmodel文件，你只需要定义一个空的加载逻辑（因为你不需要使用load_model）：

具体代码 #

import os
import torch
import mlflow.pyfunc

class DummyPyFuncModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # 无需加载模型，空实现即可
        pass

with mlflow.start_run() as run:
    run_id = run.info.run_id
    model_dir = "pyfunc_model"
    os.makedirs(model_dir, exist_ok=True)
    
    # 保存state_dict到模型目录
    state_dict_path = os.path.join(model_dir, "model_state.pth")
    torch.save(cli.model.state_dict(), state_dict_path)
    
    # 用pyfunc log模型，自动生成MLmodel文件
    mlflow.pyfunc.log_model(
        artifact_path="pyfunc_model",
        python_model=DummyPyFuncModel(),
        artifacts={"state_dict": "model_state.pth"}
    )
    
    # 注册模型
    model_uri = f"runs:/{run_id}/pyfunc_model"
    mlflow.register_model(model_uri, "Test")

完成注册后，你可以用原来的逻辑验证：

from mlflow.tracking import MlflowClient

client = MlflowClient()
run_id = client.get_model_version_by_alias("Test", "your_alias").run_id
checkpoint_artifacts = client.list_artifacts(run_id, "pyfunc_model")  # 替换为你的artifact路径

此时模型注册表会生成对应的版本，且能正常获取run_id和artifact列表。

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