MLFlow无法记录PyTorch模型：yaml.RepresenterError报错解决方案咨询

Let's break down what's happening here and how to fix it. That error about 'cannot represent an object', '1.11.0+cu102' almost always stems from YAML failing to serialize PyTorch's version string (the one with the CUDA suffix) when MLFlow tries to log model metadata. The version string might be stored as a special internal object instead of a plain string, which trips up the YAML serializer.

Here are three actionable solutions to try, ordered by simplicity:

1. Force PyTorch Version to Be a Plain String #

The quickest fix is to explicitly convert PyTorch's version to a regular string right before logging the model. This bypasses the special object that's causing the serialization issue:

import torch
import mlflow.pytorch
import pickle

# Your existing code to move model to CPU
model = model.to("cpu")
assert next(model.parameters()).device.type == "cpu"

# Add this line to fix the version serialization
torch.__version__ = str(torch.__version__)

# Now log the model as usual
mlflow.pytorch.log_model(model, artifact_path="model", pickle_module=pickle)

2. Wrap the Model with MLFlow PyFunc #

If the first fix doesn't work, you can bypass MLFlow's default PyTorch logging logic by wrapping your model as a PyFunc. This gives you more control over serialization:

import torch
import mlflow.pyfunc
import sys
import pandas as pd

class TorchModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self, model):
        self.model = model.eval()  # Set model to evaluation mode
    
    def predict(self, context, model_input):
        # Adjust this logic to match your model's input/output format
        # Example assumes input is a pandas DataFrame, tweak as needed
        with torch.no_grad():
            input_tensor = torch.tensor(model_input.values, dtype=torch.float32)
            outputs = self.model(input_tensor)
            return outputs.numpy()

# Wrap your CPU model
wrapped_model = TorchModelWrapper(model)

# Log the wrapped model with explicit conda environment
mlflow.pyfunc.log_model(
    artifact_path="model",
    python_model=wrapped_model,
    conda_env={
        "name": "mlflow-torch-env",
        "channels": ["conda-forge"],
        "dependencies": [
            f"python={'.'.join(map(str, sys.version_info[:3]))}",
            "pip",
            {"pip": [
                "torch==1.11.0+cu102",
                "mlflow",
                "pandas"
            ]}
        ]
    }
)

Make sure to adjust the predict method and conda dependencies to match your model's specific input/output requirements.

3. Upgrade MLFlow and PyTorch #

This error is a known compatibility issue between older versions of MLFlow and PyTorch (especially those with CUDA suffixes). Upgrading MLFlow to v2.0+ and ensuring your PyTorch version aligns with MLFlow's supported versions can resolve the serialization bug out of the box.

After trying any of these solutions, check your MLFlow artifacts to confirm the model is logged successfully—look for the model/ directory with all expected files like model.pth and conda.yaml.

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