MLFlow 3.3.2调用log_model()无法生成模型输入输出Schema问题

核心问题排查与解决步骤 #

1. 输入张量设备问题 #

如果输入示例张量在GPU上，直接调用.numpy()会导致隐性序列化问题（无报错但Schema无法生成），必须先转移到CPU：

inputs_example = inputs_example.cpu().numpy()

2. 正确生成Signature #

infer_signature需要同时拿到输入和输出的有效示例，且模型必须处于eval模式（避免batchnorm/dropout干扰输出结果）：

import mlflow
from mlflow.models import infer_signature
import torch

# 切换模型到评估模式
model.eval()

# 获取输入示例并转CPU+Numpy格式
dataiter = iter(train_dataloader)
inputs_example, labels_example = next(dataiter)
inputs_example = inputs_example.cpu().numpy()

# 生成输出示例（禁用梯度计算减少开销）
with torch.no_grad():
    # 从Numpy转回张量，确保设备与模型匹配
    input_tensor = torch.tensor(inputs_example).to(model.device)
    outputs_example = model(input_tensor).cpu().numpy()

# 生成完整的输入输出签名
signature = infer_signature(inputs_example, outputs_example)

3. 修正log_model参数调用 #

MLFlow 3.x中name参数已废弃，改用registered_model_name注册模型，artifact_path是必填的模型存储路径，需同时传入signature和input_example：

mlflow.pytorch.log_model(
    pytorch_model=model,
    artifact_path="srcnn_model",  # 必填：模型在工件库中的存储路径
    signature=signature,
    input_example=inputs_example,
    registered_model_name=model_name  # 可选：注册到模型仓库的名称
)

4. 验证输入维度匹配 #

确保输入示例的维度完全符合模型预期（比如SRCNN通常是(batch, channels, H, W)格式），维度不匹配会导致MLFlow无法自动推断Schema。

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