如何从run_id的mlflow中获取模型

Question

我试图检索保存在mlflow模型注册表中的pytorch模型，但未能弄清楚如何准确地这样做。我设法通过过滤所有实验的experiment_name来获得run_id，然后按我喜欢的度量进行排序，并在结果数据框中获取第一个实验的run_id。

然而，我现在停下来了。在mlflow文档中，我找不到任何方便的函数来帮助我从run_id(它应该是唯一的，特定于单个模型运行与特定的参数和参数集)模型对象过渡。我知道mlflow.pytorch.load_model()也可以接受运行URI，但我不明白为什么我需要将"run-relative/path/"传递给"artifact"。分别与runs:/<run_id>相连;此外，如何获得正确的运行URI。

我当前的代码如下:

import mlflow
def get_best_run_id(experiment_name, metric_to_sort_by='f1_score'):

experiments = mlflow.list_experiments()
experiment_path = [exp.name for exp in experiments if experiment_name in exp.name][0]

experiment_dict = dict(mlflow.get_experiment_by_name(experiment_path))

experiment_id = experiment_dict['experiment_id']
runs_df = mlflow.search_runs(experiment_id)

runs_df_filtered = runs_df[runs_df['status'] == 'FINISHED'][[col for col in runs_df.columns if col.startswith('params.') or col.startswith('metrics.')]].dropna()
runs_df_indices = runs_df_filtered.sort_values(f"metrics.{metric_to_sort_by}", ascending=False).index

runs_df = runs_df.loc[runs_df_indices, :]  
best_run_id = runs_df['run_id'].tolist()[0]

return best_run_id

best_run_id = get_best_run_id("<my_experiment_name>") 
mlflow.pytorch.load_model(f"runs:/{best_run_id}")

Answer 1

这与您如何记录您的模型有关。有几种定义model_uri的方法。如文档所示:

# Log the model
with mlflow.start_run() as run:
mlflow.pytorch.log_model(model, "model")
# Inference after loading the logged model
model_uri = "runs:/{}/model".format(run.info.run_id)
loaded_model = mlflow.pytorch.load_model(model_uri)

代码将模型记录为"model"。然后，在加载模型时，它会插入"model";URI中。您的代码没有显示您之前如何记录模型。也许，你应该回去编辑它。