MLflow and neuralforecast

import logging
import os
import warnings

import matplotlib.pyplot as plt
import mlflow
import mlflow.data
import numpy as np
import pandas as pd
from mlflow.client import MlflowClient
from mlflow.data.pandas_dataset import PandasDataset
from utilsforecast.plotting import plot_series

from neuralforecast.core import NeuralForecast
from neuralforecast.models import NBEATSx
from neuralforecast.utils import AirPassengersDF
from neuralforecast.losses.pytorch import MAE

os.environ['NIXTLA_ID_AS_COL'] = '1'
logging.getLogger("mlflow").setLevel(logging.ERROR)
logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
warnings.filterwarnings("ignore")

# Split data and declare panel dataset
Y_df = AirPassengersDF
Y_train_df = Y_df[Y_df.ds<='1959-12-31'] # 132 train
Y_test_df = Y_df[Y_df.ds>'1959-12-31'] # 12 test
Y_df.tail()

mlflow.pytorch.autolog(checkpoint=False)

with mlflow.start_run() as run:
    # Log the dataset to the MLflow Run. Specify the "training" context to indicate that the
    # dataset is used for model training
    dataset: PandasDataset = mlflow.data.from_pandas(Y_df, source="AirPassengersDF")
    mlflow.log_input(dataset, context="training")

    # Define and log parameters
    horizon = len(Y_test_df)
    model_params = dict(
        input_size=1 * horizon,
        h=horizon,
        max_steps=300,  
        loss=MAE(),
        valid_loss=MAE(),  
        activation='ReLU',
        scaler_type='robust',
        random_seed=42,
        enable_progress_bar=False,
    )
    mlflow.log_params(model_params)

    # Fit NBEATSx model
    models = [NBEATSx(**model_params)]
    nf = NeuralForecast(models=models, freq='M')           
    train = nf.fit(df=Y_train_df, val_size=horizon)
    
    # Save conda environment used to run the model
    mlflow.pytorch.get_default_conda_env()
    
    # Save pip requirements
    mlflow.pytorch.get_default_pip_requirements()

mlflow.pytorch.autolog(disable=True)

# Save the neural forecast model
nf.save(path='./checkpoints/test_run_1/',
        model_index=None, 
        overwrite=True,
        save_dataset=True)

Seed set to 42

Y_hat_df = nf.predict(futr_df=Y_test_df)
plot_series(Y_train_df, Y_hat_df, palette='tab20b')