Controlling the Number of CPU Cores Used in N-Beats Model Training with Darts

243 Views Asked by At

I am working with the Darts library to train an N-Beats model in a JupyterLab environment running on a Linux virtual machine, accessed via Mobaxterm. I have set up the PyTorch Lightning trainer (used internally by Darts) to use the CPU and specified one device for training. However, when I monitor the virtual machine via Mobaxterm, I see that it is using four CPU cores instead of the expected two.

Here is the relevant part of my code:

pl_trainer_kwargs = {"accelerator": "cpu", "devices" : 2, "callbacks":[loss_logger]}

model_cov_2 = NBEATSModel(
    # model parameters
    pl_trainer_kwargs=pl_trainer_kwargs
)

model_cov_2.fit(
    # fit parameters
) 

Is there a way to limit the number of CPU cores used by the model during training to just two cores in this Linux environment? Any help or guidance would be appreciated.

I've tried to limit the number of CPU cores used by PyTorch by setting the OMP_NUM_THREADS environment variable to "2" as follows:

import os
os.environ["OMP_NUM_THREADS"] = "2"
0

There are 0 best solutions below