MLFlow无法记录pytorch模型

Question

在训练一个模型之后，我尝试用mlflow.pytorch.log_model(model, artifact_path="model",pickle_module=pickle)记录它到mlflow，但是我得到了错误：

yaml.representer.RepresenterError: ('cannot represent an object', '1.11.0+cu102')

在此之前，我肯定会将模型发送到cpu，并使用next(model.parameters()).device确认其存在。如何获得mlflow来记录我的模型？

Answer 1

我没有在MLflow中使用PyTorch，但在最坏的情况下，您可以避免使用PyTorch风格，只需执行以下操作

    with tempfile.TemporaryDirectory() as path:
        model.save(path)
        mlflow.log_artifacts(path, artifact_path="model")