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")