Demux (fork) can not be serialized if there is a child not depleted

[ejguan]

🐛 Describe the bug

As shown in pytorch/vision#5187 (comment)

Not just special use case like vision. If users use in_memory_cache or any data structure to capture/store the data from any ChildDataPipe from demux or fork, the pipeline becomes non-serializable from the second epoch.

This may not interfere the DataLoader iteration. But, it is definitely going to block snapshotting.

cc: @NivekT Do you plan to work on this?

Versions

PyTorch core master

[NivekT]

Yes, I am aware. I believe any DataPipes that contain a buffer will not be serializable without custom __getstate__ and __setstate__ methods that avoid storing any components that are not serializable (e.g. the buffer).

If I understand correctly, this is not blocking anything in torchvision, correct?

[ejguan]

If I understand correctly, this is not blocking anything in torchvision, correct?

TorchVision is actually blocked by this. They can't use demux if DataLoader is involved. The another workaround Phillip proposed would also not work. We have to fix it.

[ejguan]

Yes, I am aware. I believe any DataPipes that contain a buffer will not be serializable without custom getstate and setstate methods that avoid storing any components that are not serializable (e.g. the buffer).

If buffer is attached to instance, serialization won't work. If buffer is created within iterator, then it's fine.

[ejguan]

I have no idea why I created such issue, as I have sent a PR to fix it. pytorch/pytorch#71311
Closing. Sorry about that @NivekT

[ejguan]

Actually Fork is not fixed...