Mlflow log_figure删除神器

Question

我正在运行带有自动记录功能的 mlflow 来跟踪 xgboost 模型。默认情况下，它在工件下保存模型、需求和功能重要性。我想保留很酷的东西。

但是，如果我尝试使用下面的代码添加图形，它会删除工件文件夹并只将我的图形放入其中。将工件附加到文件夹的最佳方法是什么？

mlflow.autolog()
mlflow.set_experiment('Energy Use Forecasting')

def main():
    color_pal = sns.color_palette()
    input_file = Path(r'D:\data\ML\PowerConsumption\AEP_hourly.csv')
    models_path = Path(r'D:\Models\ML') / Path(__file__).stem
    models_path.mkdir(parents=True, exist_ok=True)

    df = pd.read_csv(input_file)
    df = df.set_index('Datetime')
    df.index = pd.to_datetime(df.index)

    data_source = input_file.stem.split('_')[0]
    df = df.rename(columns={data_source + '_MW': 'MW'})
    train, val = train_val_split(df)

    model = xgb.XGBRegressor(base_score=0.5, booster='gbtree',
                                 n_estimators=10000,
                                 early_stopping_rounds=50,
                                 objective='reg:squarederror',
                                 max_depth=max_depth,
                                 learning_rate=learning_rate)
    model.fit(X_train, y_train,
                  eval_set=[(X_train, y_train), (X_val, y_val)],
                  verbose=100)

    ax = df[['MW']].plot(figsize=(15, 5))
    df['prediction'].plot(ax=ax, style='.')
    plt.legend(['Truth Data', 'Predictions'])
    plt.axvline(val_split_index, color="gray", lw=3, label=f"Val split point")
    ax.set_title('Raw Data and Prediction')

    mlflow.log_figure(fig_trainval_preds, 'trainval_predictions.png')

Answer 1

问题在于 mlflow.log_figure() 不知道 run_id 并创建一个新的运行来登录。尽管在 mlflow.autolog() 使用的一些示例中未显示，但需要使用以下内容包装任何模型拟合和绘图代码：

with mlflow.start_run() as run:

然后，手动日志记录将转到与自动日志记录相同的 run_id。