After some offline discussion with @NicolasHug, it became clear that the situation is not as straight forward as I thought it was. For example, the child datapipes from the Demultiplexer seem to be aware if an iteration is ongoing.

@ejguan @NivekT @VitalyFedyunin what is the correct way here?

IMHO, it depends on your Dataset:

1. If you want to pre-load labels or all meta data before iteration, I think it might be better to use filter for label_dp and image_dp and deplete label_dp before iteration starts:

label_dp = resource_dp.filter(is_meta_file)
labels = list(label_dp)

image_dp = resource_dp.filter(is_image_file)

2. If meta data or labels can be loaded lazily, IterToMap might be a better solution here. The meta data is going to be loaded lazily. And, zip between image_dp and label_dp.
   • to_map_datapipe: https://github.com/pytorch/data/blob/13b574c80e8732744fee6ab9cb7e35b5afc34a3c/torchdata/datapipes/iter/util/converter.py#L21
   • zip_with_map: https://github.com/pytorch/data/blob/13b574c80e8732744fee6ab9cb7e35b5afc34a3c/torchdata/datapipes/iter/util/combining.py#L148-L149

label_dp, image_dp = resource_dp.demux(classify_fn, 2)

label_dict = label_dp. to_map_datapipe()
image_dp.zip_with_map(label_dict)
...

After some offline discussion with @NicolasHug, it became clear that the situation is not as straight forward as I thought it was. For example, the child datapipes from the Demultiplexer seem to be aware if an iteration is ongoing.>

That only matters if you have read one of your child DataPipes eagerly during the construction (it seems that is no longer the case after your code change in this PR since it no longer calls demux in cases where labels are eagerly needed).

So for my own understanding: the result of IterToMapConverter is something similar to what I proposed in #5219, right? Basically we get a dictionary that will only be filled at runtime, correct?

@ejguan's proposal was to use the map together with a MapKeyZipper. While playing with that I found that it is quite inconvenient since we cannot use a lambda or local function for the merge_fn. I preferred to use the map directly with

dp = ...
map = IterToMapConverter(dp)

other_dp = ...
other_dp = Mapper(other_dp, map.__getitem__, input_col=...)

Do you see any downside with that?

Regardless of the approach chosen, I found that using IterToMapConverter introduces a performance degradation. Take the following snippet:

from time import perf_counter

import torch
from torchvision.prototype import datasets

for name, config in [
    ("imagenet", dict(split="val")),
    ("cub200", dict(year="2011")),
]:
    warmup_times = []
    for _ in range(5):
        tic = perf_counter()
        for sample in datasets.load(name, **config):
            break
        tac = perf_counter()

        warmup_times.append(tac - tic)

    print(f"Warmup for {name} took on average {float(torch.tensor(warmup_times).mean()):.2f} seconds")

Running this on main prints

Warmup for imagenet took on average 2.53 seconds
Warmup for cub200 took on average 2.97 seconds

Running it on this PR in the current state prints

Warmup for imagenet took on average 2.77 seconds
Warmup for cub200 took on average 3.64 seconds

Any idea what could cause this? I'm aware that meta-pytorch/data#454 proposes a speed-up, but the implementation on main currently also loads the full dictionary at once. Thus, this is not the cause here.

@pmeier you might be interested in #6128 which IIRC attemps to fix this issue as well

Is the issue related to ResourceWarning: unclosed (which only happens for Python 3.9 Windows) caused by what you described above, or is that separate? I think #6128 is trying to fix that.

Yes, this is related. Although #6128 patches more things, I think this PR is sufficient to get rid of the warnings, given that they were always related to items left in a Demultiplexer.