add DataLoader related parameters to fit() and predict()

[BohdanBilonoh]

Checklist before merging this PR:

• Mentioned all issues that this PR fixes or addresses.
• Summarized the updates of this PR under Summary.
• Added an entry under Unreleased in the Changelog.

Summary

Add torch.utils.data.DataLoader related parameters to fit() and predict() of TorchForecastingModel

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 93.74%. Comparing base (a0cc279) to head (ec7a01b).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2295      +/-   ##
==========================================
- Coverage   93.75%   93.74%   -0.01%     
==========================================
  Files         138      138              
  Lines       14352    14346       -6     
==========================================
- Hits        13456    13449       -7     
- Misses        896      897       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[madtoinou]

Hi @BohdanBilonoh,

It looks great, however to make it easier to maintain and more exhaustive, I think that it would be great to just add an argument called dataloader_kwargs, then check that the argument explicitly used by Darts are not redundant/overwritten and then pass this argument down to the DataLoader constructor.

It will allow users to specify more than just prefetch_factor, persistent_workers and pin_memory, while limiting copy-pasting from other library documentation (putting a link to the torch.DataLoader page does sound like a good idea for this argument however)

PS: Apologies for taking so long with the review of this PR.

[joshua-xia]

Hi @BohdanBilonoh, Would you please add multiprocessing_context parameter for Dataloader, it is useful when we use multi-workers for dataloader, Thanks!

[joshua-xia]

Hi @BohdanBilonoh, Would you please add multiprocessing_context parameter for Dataloader, it is useful when we use multi-workers for dataloader, Thanks!

@BohdanBilonoh refer to #2375

[joshua-xia]

Hi @BohdanBilonoh, Would you please add multiprocessing_context parameter for Dataloader, it is useful when we use multi-workers for dataloader, Thanks!

@BohdanBilonoh refer to #2375

@BohdanBilonoh My bad, it is good idea from @madtoinou to add dataloader_kwargs to let user input dataloader parameters as wish freely, not need to support special multiprocessing_context parameter forcibly

[BohdanBilonoh]

@madtoinou what do you think about hardcoded parameters like

batch_size=self.batch_size,
shuffle=True,
drop_last=False,
collate_fn=self._batch_collate_fn,

should it be hard coded for new dataloader_kwargs like

def _setup_for_train(
    self,
    train_dataset: TrainingDataset,
    val_dataset: Optional[TrainingDataset] = None,
    trainer: Optional[pl.Trainer] = None,
    verbose: Optional[bool] = None,
    epochs: int = 0,
    dataloader_kwargs: Optional[Dict[str, Any]] = None,
) -> Tuple[pl.Trainer, PLForecastingModule, DataLoader, Optional[DataLoader]]:

        ...

        if dataloader_kwargs is None:
            dataloader_kwargs = {}

        dataloader_kwargs["shuffle"] = True
        dataloader_kwargs["batch_size"] = self.batch_size
        dataloader_kwargs["drop_last"] = False
        dataloader_kwargs["collate_fn"] = self._batch_collate_fn

        # Setting drop_last to False makes the model see each sample at least once, and guarantee the presence of at
        # least one batch no matter the chosen batch size
        train_loader = DataLoader(
            train_dataset,
            **dataloader_kwargs,
        )

        dataloader_kwargs["shuffle"] = False

        # Prepare validation data
        val_loader = (
            None
            if val_dataset is None
            else DataLoader(
                val_dataset,
                **dataloader_kwargs,
            )
        )

        ...

or give a user full control on dataloader_kwargs?

[tRosenflanz]

@madtoinou what do you think about hardcoded parameters like

batch_size=self.batch_size,
shuffle=True,
drop_last=False,
collate_fn=self._batch_collate_fn,

should it be hard coded for new dataloader_kwargs like

def _setup_for_train(
    self,
    train_dataset: TrainingDataset,
    val_dataset: Optional[TrainingDataset] = None,
    trainer: Optional[pl.Trainer] = None,
    verbose: Optional[bool] = None,
    epochs: int = 0,
    dataloader_kwargs: Optional[Dict[str, Any]] = None,
) -> Tuple[pl.Trainer, PLForecastingModule, DataLoader, Optional[DataLoader]]:

        ...

        if dataloader_kwargs is None:
            dataloader_kwargs = {}

        dataloader_kwargs["shuffle"] = True
        dataloader_kwargs["batch_size"] = self.batch_size
        dataloader_kwargs["drop_last"] = False
        dataloader_kwargs["collate_fn"] = self._batch_collate_fn

        # Setting drop_last to False makes the model see each sample at least once, and guarantee the presence of at
        # least one batch no matter the chosen batch size
        train_loader = DataLoader(
            train_dataset,
            **dataloader_kwargs,
        )

        dataloader_kwargs["shuffle"] = False

        # Prepare validation data
        val_loader = (
            None
            if val_dataset is None
            else DataLoader(
                val_dataset,
                **dataloader_kwargs,
            )
        )

        ...

or give a user full control on dataloader_kwargs?

you could extend your suggestion to allow overrides but with populated defaults

defaults = dict(shuffle = True, batch_size = self.batch_size, drop_last = False, collate_fn = self._batch_collate_fn)
#combine with defaults but override them
dataloader_kwargs_train = {**defaults,**dataloader_kwargs}
#override shuffle
dataloader_kwargs_val = (**dataloader_kwargs_train, **dict(shuffle=False)} 

[dennisbader]

Thanks for the updates @BohdanBilonoh, I took the chance to fix some indentation issues in one of the docstrings and pushed the changes.

If we give the freedom to overwrite our default dataloader params when calling fit(), shouldn't we then also allow that during predict()?

Also, removing the num_loader_workers parameter is a breaking change. Can you document this in the CHANGELOG.md?

[dennisbader]

Shouldn't we allow then here as well the liberty to overwrite these defaults (expect shuffle)?

[dennisbader]

it's referring to the train and val datasets, but it should be for the prediction dataset, same for the methods below