DataModules: add configurable args to dataloader

[anton-emanuel]

Closes #2332.

[adamjstewart]

How much of a speed difference do you notice?

At the moment, the datamodules are all designed to pass additional kwargs to the dataset class. So we should technically add these new args to every datamodule, not just to the base class. We could design things such that all kwargs are either passed to the data loader or the dataset class, but I'm worried users might get confused by this. Of course, they don't need to understand the feature if they don't want to use it.

I guess my concern is, why stop at only these variables?

[anton-emanuel]

Varying all of those three parameters I was able to at least half the time for [_TrainingEpochLoop].train_dataloader_next (probably even less). But I am reading from COGs over internet, so I guess it's very different from having the raster data on the same device.

I agree, it would be better to solve for all arguments, and that there is a trade off between full configurability and simplicity. I just added those variables since those were the ones I wanted to vary.

Could it be an idea to do something like what is done in BaseDataModule._valid_attribute, but just taking kwargs starting with dataloader_ and put in dataloader (I feel like I have seen this solution somewhere in TorchGeo, but might have been another library)?

# ...
DATALOADER_KWARG_PREFIX = "dataloader_"

    # ...
    dataloader_kwargs = {
        k.replace(DATALOADER_KWARG_PREFIX, ""): v
        for k, v in kwargs.items()
        if k.startswith(DATALOADER_KWARG_PREFIX)
    }
    
    # ...
    DataLoader(
        dataset=dataset,
        batch_size=batch_size,
        sampler=sampler,
        batch_sampler=batch_sampler,
        num_workers=self.num_workers,
        collate_fn=self.collate_fn,
        persistent_workers=self.num_workers > 0,
        **dataloader_kwargs
    )

[anton-emanuel]

I changed it to use kwargs now

        self.dataloader_kwargs, self.kwargs = split_prefixed_kwargs(
            'dataloader_', **kwargs
        )

[anton-emanuel]

@microsoft-github-policy-service agree company="Kongsberg Satellite Services"

[DimitrisMantas]

I think it woul be better for kwargs to be a JSON-like dict, something like {"dataset": ..., "dataloader": ...}. We could even turn this into a typed dict or a dataclass so that it can be clearly documented and tested. I think the transformers library likes to use these config structures a lot