DataLoader wrapping, re-instantiation and patching

[awaelchli]

🐛 Bug

Recent changes to the Lite DataLoader in #10279 have not been properly reviewed and tested. Several clean ups need to be done and better tests need to be written.

The following issue exists:

To Reproduce

def test_setup_custom_dataloaders_new():

    class DataLoaderSubclass1(DataLoader):
        def __init__(self, attribute1, *args, **kwargs):
            # UNCOMMENT THE LINE BELOW, NO DIFFERENCE (EXPECTED)
            # self.attribute1 = attribute1
            super().__init__(*args, **kwargs)

    class DataLoaderSubclass2(DataLoaderSubclass1):
        def __init__(self, attribute1, attribute2, *args, **kwargs):
            # UNCOMMENT THE LINE BELOW, TEST PASSES
            # self.attribute2 = attribute2
            super().__init__(attribute1, *args, **kwargs)

    class LiteWithCustomDataLoader(LightningLite):
        def run(self):
            dataloader = DataLoaderSubclass2("attribute1", "attribute2", dataset=range(4), batch_size=2)
            lite_dataloader = self.setup_dataloaders(dataloader)
            
            # THIS ASSERTION PASSES
            assert lite_dataloader._dataloader.attribute1 == "attribute1"
            
            # THIS ASSERTION FAILS
            assert lite_dataloader._dataloader.attribute2 == "attribute2"

    LiteWithCustomDataLoader().run()

Expected behavior

According to the changes in the PR #10279, this is supposed to work and was the main motivation of it. But it does not because the tests were insufficient. This issue was found by me in #10297 while trying to clean up the code.

Environment

Lightning v1.5.0!

cc @tchaton @rohitgr7 @carmocca @justusschock @awaelchli

[awaelchli]

The issue exists because this wrapper here

https://github.com/PyTorchLightning/pytorch-lightning/blob/93caa7cda9277f86f17b1eab0793019aee8d40c2/pytorch_lightning/lite/wrappers.py#L106-L116

pops args and kwargs.
When we get to this line here

https://github.com/PyTorchLightning/pytorch-lightning/blob/93caa7cda9277f86f17b1eab0793019aee8d40c2/pytorch_lightning/lite/lite.py#L242-L244

we are re-instantiating (with the new init) but passing arguments by kwargs, strictly! This means the for loop here runs empty!

https://github.com/PyTorchLightning/pytorch-lightning/blob/93caa7cda9277f86f17b1eab0793019aee8d40c2/pytorch_lightning/lite/wrappers.py#L110-L114

What is the design flaw? As always, it is the fact that we are making too strong assumptions without documenting them and without testing for them.

[awaelchli]

A few things that need to be addressed (discussed with @tchaton and @carmocca):

• Bugfix PR Fix DataLoader inspection and re-instantiation in Lite #10334
• Revert dataset TypeError Revert part of #10279 #10376
  https://github.com/PyTorchLightning/pytorch-lightning/blob/1c28f361d47257feb0bd34cbc8d41642a1f0a2ab/pytorch_lightning/lite/lite.py#L242-L248
• Handle TypeError and raise a MisconfigurationError Improve error message on TypeError during DataLoader reconstruction #10719
• Also use the the context manager around dataloader methods in Lightning Support re-instantiation for custom DataLoader in Lightning #10680
  https://github.com/PyTorchLightning/pytorch-lightning/blob/1c28f361d47257feb0bd34cbc8d41642a1f0a2ab/pytorch_lightning/lite/lite.py#L408-L410
• Open PR to PyTorch Geometric to improve support for their DataLoaders

[awaelchli]

One minor limitation of the patching is that the dataloader needs to allow write access to its dict, i.e. something like this would be a valid DataLoader but our patching would fail:

class MyDataLoader(DataLoader):

    def __init__(self, whatever, *args, **kwargs):
        # Lightning will not be allowed to set the attribute self.whatever here (read-only)
        super().__init__(*args, **kwargs)
        
    @property
    def whatever(self):
        return "foo"

We may want to handle this error case.

[carmocca]

In practice, it would be instead:

class MyDataLoader(DataLoader):

    def __init__(self, whatever, *args, **kwargs):
        self._whatever = whatever
        super().__init__(*args, **kwargs)
        
    @property
    def whatever(self):
        return self._whatever

[awaelchli]

@carmocca @tchaton Do you remember what we wanted to do with "Open PR to PyTorch Geometric to improve support for their DataLoaders"?
Was it about removing these lines?
https://github.com/pyg-team/pytorch_geometric/blob/4bd1821b72039cb6beb96041bc9f18b2c4af36bb/torch_geometric/loader/dataloader.py#L74-L76

But we need to wait until at least 1.6 is out.

[carmocca]

No, it was about addressing the failure which made @tchaton add the TypeError catch. But I'm not sure which implementation caused the error, most already drop the arguments from kwargs.

[tchaton]

Hey @awaelchli,

I wonder if it would be this one: https://github.com/pyg-team/pytorch_geometric/blob/4bd1821b72039cb6beb96041bc9f18b2c4af36bb/torch_geometric/loader/graph_saint.py#L50 where the dataset isn't being poped from the kwargs, but this is maybe expected.

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!