Transforms are lost when using a dataloader with spawn

[RickLuiken]

🐛 Describe the bug

When using a DataLoader with multiprocessing_context="spawn", transforms are lost. I believe this is an issue with the pickle change in #7860, the transforms are not pickled.

Minimal example:

from torch.utils.data import DataLoader
from torchvision.transforms import v2 as T
from torchvision.datasets import VOCDetection, wrap_dataset_for_transforms_v2

def _no_collate(batch):
    return batch

transform = T.ToImage()
dataset = VOCDetection(..., transforms=transform)
wrapped_dataset = wrap_dataset_for_transforms_v2(dataset)

dataloader = DataLoader(wrapped_dataset, num_workers=2, multiprocessing_context="spawn", collate_fn=_no_collate)

# Sample from the dataset is not a tv_tensor image!
next(iter(dataloader))

Versions

Collecting environment information...
PyTorch version: 2.1.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Enterprise
GCC version: (Rev10, Built by MSYS2 project) 12.2.0
Clang version: Could not collect
CMake version: version 3.26.0
Libc version: N/A

Python version: 3.11.3 (tags/v3.11.3:f3909b8, Apr 4 2023, 23:49:59) [MSC v.1934 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22621-SP0
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX A2000 Laptop GPU
Nvidia driver version: 536.25
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=2304
DeviceID=CPU0
Family=198
L2CacheSize=10240
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=2304
Name=11th Gen Intel(R) Core(TM) i7-11800H @ 2.30GHz
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] flake8==6.1.0
[pip3] mypy==1.6.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.3
[pip3] torch==2.1.0+cu121
[pip3] torchaudio==2.1.0+cu121
[pip3] torchvision==0.16.0+cu121
[conda] Could not collect

cc @vfdev-5

[NicolasHug]

Thanks for the report @RickLuiken , I can reproduce. @pmeier is this something you'll have some cycles to look into? LMK if not

[pmeier]

Can do but likely not before EOW. When is our bug fix release planned. Can we maybe get this in?

[NicolasHug]

Yeah we definitely should try to get it before the bugfix release. The cherry-pick deadline is Nov 3 (release aimed for mid Nov).

I started looking into it on my side, the problem is related to

vision/torchvision/tv_tensors/_dataset_wrapper.py

Line 174 in 3fb88b3

self.transform, dataset.transform = dataset.transform, None

and to the fact that with spawn (not with fork), this line in __init__ is called num_workers times (EDIT: it's actually num_workers + 1 times, see below). And starting from the second time, dataset.transform is None and self.transform ends up being None too, leading to no transformation.

[pmeier]

Shouldn't each spawned process have its own copy of the dataset?

[NicolasHug]

Yeah they do. But all num_workers processes are spawned from the main/parent process, passing through __reduce__ which calls

vision/torchvision/tv_tensors/_dataset_wrapper.py

Line 202 in 3fb88b3

return wrap_dataset_for_transforms_v2, (self._dataset, self._target_keys)

at which point the self._dataset's [target_]transform[s] fields were already set to None in at the line mentioned above #8066 (comment). So we end up with num_workers new dataset objects which have no transforms.

Here's a sequence of event which will help clarify:

first number is pid. num_workers = 4

799599 In wrap_dataset_for_transforms_v2, dataset.transform = Compose(    ToImage())
799599 In VisionDatasetTVTensorWrapper.__init__, self.transform = Compose(    ToImage()), dataset.transform = None
799599 In __reduce__
799599 In __reduce__
799599 In __reduce__
799599 In __reduce__
799601 In wrap_dataset_for_transforms_v2, dataset.transform = None
799601 In VisionDatasetTVTensorWrapper.__init__, self.transform = None, dataset.transform = None
799601 In VisionDatasetTVTensorWrapper.__getitem__, self.transform = None
799602 In wrap_dataset_for_transforms_v2, dataset.transform = None
799602 In VisionDatasetTVTensorWrapper.__init__, self.transform = None, dataset.transform = None
799602 In VisionDatasetTVTensorWrapper.__getitem__, self.transform = None
799603 In wrap_dataset_for_transforms_v2, dataset.transform = None
799603 In VisionDatasetTVTensorWrapper.__init__, self.transform = None, dataset.transform = None
799603 In VisionDatasetTVTensorWrapper.__getitem__, self.transform = None
799604 In wrap_dataset_for_transforms_v2, dataset.transform = None
799604 In VisionDatasetTVTensorWrapper.__init__, self.transform = None, dataset.transform = None
799604 In VisionDatasetTVTensorWrapper.__getitem__, self.transform = None