nbeats model logging in mlflow not working

Question

I am trying to log nbeats model developed using darts library into mlflow ui, but I am not able to do that. I am getting the following error.

PS this model is developed using Darts library. If anyone knows who had developed nbetas and logged in mlflow. Kindly mention the source for my reference.

RuntimeError: [enforce fail at C:\cb\pytorch_1000000000000\work\c10\core\impl\alloc_cpu.cpp:81] data. DefaultCPUAllocator: not enough memory: you tried to allocate 16515072 bytes.

My code is as folows:

from darts.models import LightGBMModel, RandomForest, NaiveMean, NBEATSModel
from data_preparation import darts_x_train, darts_y_train, darts_y_val, df_x

import pandas as pd
import tensorflow as tf
import mlflow
import warnings
warnings.filterwarnings("ignore")
import torch

print(torch.__version__)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

import logging

logging.basicConfig(level=logging.WARN)
logger = logging.getLogger(__name__)
lags = [-4, -24 * 4, -12 * 4, -8 * 4 * 24, -7 * 24 * 4, -10 * 4 * 24, -11 * 4 * 24, -12 * 4 * 24]
lags_future = [0, 4, 8, 12 * 4, 24 * 4, 48 * 4]
output_chunk_length = 4  # int(pd.Timedelta('2w') / (darts_x_train.time_index[1] - 
darts_x_train.time_index[0]))
print(output_chunk_length, "output_chunk_length timesteps")
last_two_weeks = darts_y_val.time_index[-1] - pd.Timedelta(days=16)
test_series = darts_y_val.drop_after(last_two_weeks)

def baseline_model():
    nbeats_model = NBEATSModel(input_chunk_length=672, output_chunk_length=output_chunk_length, n_epochs=1,
                           random_state=42,
                           pl_trainer_kwargs={
                               "accelerator": "gpu",
                               "devices": [0]
                           }, )

return nbeats_model

def mlflow_run(run_name="nbeats_model_run"):

    model_uri = mlflow.get_artifact_uri("model")
    # Start MLflow run and log everyting...
    with mlflow.start_run(run_name=run_name, nested=True) as run:
    model = baseline_model()
    mlflow.keras.log_model(keras_model=model, keras_module= tf.keras, artifact_path='artifacts')
    model.fit(darts_y_train, past_covariates=darts_x_train, epochs=1)
    model.predict(n=1344, series=test_series, past_covariates=df_x)
    run_id = run.info.run_uuid
    exp_id = run.info.experiment_id

    return exp_id, run_id


 if __name__ == '__main__':
# suppress any deprecated warnings
 warnings.filterwarnings("ignore", category=DeprecationWarning)

 (exp_id, run_id) = mlflow_run()
 print(f"Finished Experiment id={exp_id} and run id = {run_id}")

C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\python.exe "C:/Users/R953682/Local_Energy _Consumption/MLW/run_file.py" 34397 34397 11467 11467 1.12.1 4 output_chunk_length timesteps 2022/09/22 17:31:09 WARNING mlflow.utils.environment: Encountered an unexpected error while inferring pip requirements (model URI: C:\Users\R953682\AppData\Local\Temp mpnefahsun\model, flavor: keras), fall back to return ['tensorflow==2.9.1', 'keras==2.9.0']. Set logging level to DEBUG to see the full traceback. 2022-09-22 17:31:10 pytorch_lightning.utilities.rank_zero INFO: GPU available: True (cuda), used: True 2022-09-22 17:31:10 pytorch_lightning.utilities.rank_zero INFO: TPU available: False, using: 0 TPU cores 2022-09-22 17:31:10 pytorch_lightning.utilities.rank_zero INFO: IPU available: False, using: 0 IPUs 2022-09-22 17:31:10 pytorch_lightning.utilities.rank_zero INFO: HPU available: False, using: 0 HPUs 2022-09-22 17:31:10 pytorch_lightning.accelerators.cuda INFO: LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0] 2022-09-22 17:31:11 pytorch_lightning.callbacks.model_summary INFO: | Name | Type | Params

0 | criterion | MSELoss | 0 1 | train_metrics | MetricCollection | 0 2 | val_metrics | MetricCollection | 0 3 | stacks | ModuleList | 69.4 M

69.3 M Trainable params 49.7 K Non-trainable params 69.4 M Total params 555.182 Total estimated model params size (MB) Epoch 0: 100%|██████████| 1054/1054 [02:46<00:00, 6.33it/s, loss=0.0392, train_loss=0.0496] 2022-09-22 17:33:57 pytorch_lightning.utilities.rank_zero INFO: Trainer.fit stopped: max_epochs=1 reached. Traceback (most recent call last): File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning rainer rainer.py", line 650, in _call_and_handle_interrupt return trainer_fn(*args, **kwargs) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning rainer rainer.py", line 735, in _fit_impl results = self._run(model, ckpt_path=self.ckpt_path) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning rainer rainer.py", line 1169, in _run self._teardown() File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning rainer rainer.py", line 1229, in _teardown self.strategy.teardown() File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning\strategies\strategy.py", line 476, in teardown self.lightning_module.cpu() File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning\core\mixins\device_dtype_mixin.py", line 141, in cpu return super().cpu() File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 738, in cpu return self._apply(lambda t: t.cpu()) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 579, in _apply module._apply(fn) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 579, in _apply module._apply(fn) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 579, in _apply module._apply(fn) [Previous line repeated 3 more times] File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 615, in _apply grad_applied = fn(param.grad) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 738, in return self._apply(lambda t: t.cpu()) RuntimeError: [enforce fail at C:\cb\pytorch_1000000000000\work\c10\core\impl\alloc_cpu.cpp:81] data. DefaultCPUAllocator: not enough memory: you tried to allocate 16515072 bytes.

During handling of the above exception, another exception occurred:

Traceback (most recent call last): File "C:\Users\R953682\Local_Energy _Consumption\MLW un_file.py", line 60, in (exp_id, run_id) = mlflow_run() File "C:\Users\R953682\Local_Energy _Consumption\MLW un_file.py", line 48, in mlflow_run model.fit(darts_y_train, past_covariates=darts_x_train, epochs=1) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\darts\utils orch.py", line 112, in decorator return decorated(self, *args, **kwargs) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\darts\models\forecasting orch_forecasting_model.py", line 806, in fit return self.fit_from_dataset( File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\darts\utils orch.py", line 112, in decorator return decorated(self, *args, **kwargs) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\darts\models\forecasting orch_forecasting_model.py", line 961, in fit_from_dataset self._train(train_loader, val_loader) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\darts\models\forecasting orch_forecasting_model.py", line 983, in _train self.trainer.fit( File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning rainer rainer.py", line 696, in fit self._call_and_handle_interrupt( File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning rainer rainer.py", line 664, in _call_and_handle_interrupt self._teardown() File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning rainer rainer.py", line 1229, in _teardown self.strategy.teardown() File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning\strategies\strategy.py", line 476, in teardown self.lightning_module.cpu() File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages\pytorch_lightning\core\mixins\device_dtype_mixin.py", line 141, in cpu return super().cpu() File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 738, in cpu return self._apply(lambda t: t.cpu()) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 579, in _apply module._apply(fn) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 579, in _apply module._apply(fn) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 579, in _apply module._apply(fn) [Previous line repeated 3 more times] File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 615, in _apply grad_applied = fn(param.grad) File "C:\Users\R953682\Anaconda3\envs\MLflow_tutorial\lib\site-packages orch n\modules\module.py", line 738, in return self._apply(lambda t: t.cpu()) RuntimeError: [enforce fail at C:\cb\pytorch_1000000000000\work\c10\core\impl\alloc_cpu.cpp:81] data. DefaultCPUAllocator: not enough memory: you tried to allocate 16515072 bytes.

nbeats model logging in mlflow not working

0 | criterion | MSELoss | 0 1 | train_metrics | MetricCollection | 0 2 | val_metrics | MetricCollection | 0 3 | stacks | ModuleList | 69.4 M

0 Answers