How to duplicate samples in a datapipe?

Vedant_Roy: <p>I have a <code>torchdata</code> data pipe which starts off with a very expensive process: reading a video into frames.</p>
<p>Ideally, I’d like to duplicate the frames from the video (when training on e.g, a single video for overfitting) to avoid reading the video from disk each time.</p>
<p>I tried using <code>.cycle</code> to do this, but that doesn’t work.<br>
An alternative option would be to just do a <code>map</code> that (for example) doubles the number of frames, but I was wondering if there is a better “recommended” solution for this type of thing.</p>

nivek: <p>Can you elaborate on what you mean by duplicating (i.e. what the input and output should look like)? And what makes <code>.cycle</code> not work for your use case?</p>
<p>Is it something like <code>[0, 1, 2] -&gt; [0, 0, 1, 1, 2, 2]</code>? In that case, I would recommend using <code>.flatmap(repeat_fn)</code>.</p>
<p>Note that if you are going to shuffle later, then <code>.cycle</code> should work.</p>

Vedant_Roy: <p><a class="mention" href="/u/nivek">@nivek</a></p>
<p>My dataloaders look like this:</p>
<pre><code class="lang-auto">def local_datapipe(dir):
    assert isinstance(dir, str), f"dir must be a str, got {type(dir)}"
    datapipe = dp.iter.FSSpecFileLister(str(dir))
    print(f"Found {len(list(datapipe))} files in {dir}")
    datapipe = datapipe.map(load_video)
    return datapipe


def load_video(path) -&gt; VideoReader:
    print("Loading video...")
    return VideoReader(path, "video")

def video_overfit_dataloader(
    loader_dp: dp.iter.IterDataPipe, fn, single_image=False, **kwargs
):
    def repeat(N, x):
        return [x for _ in range(N)]

    # `fn` turns a video into a list of frames
    dp = loader_dp.flatmap(fn)
    if single_image:
        dp = dp.header(1)
    dp = dp.flatmap(functools.partial(repeat, 5000))
    if not single_image:
        dp = dp.shuffle()
    dp = dp.sharding_filter()
    dl = DataLoader(dp, **kwargs, drop_last=True)
    assert_dl_nonempty(dl, **kwargs)
    return dl
</code></pre>
<p>Originally, I was doing:</p>
<pre><code class="lang-auto">def video_overfit_dataloader(
    loader_dp: dp.iter.IterDataPipe, fn, single_image=False, **kwargs
):
    def repeat(N, x):
        return [x for _ in range(N)]

    dp = loader_dp.flatmap(fn)
    if single_image:
        dp = dp.header(1)
        dp = dp.cycle(5000)
    if not single_image:
        dp = dp.shuffle()
    dp = dp.sharding_filter()
    dl = DataLoader(dp, **kwargs, drop_last=True)
    assert_dl_nonempty(dl, **kwargs)
    return dl
</code></pre>
<p>Note: I’m not sure if the above is exactly what I was doing originally, but I think?? it should be close.<br>
When using the <code>.cycle</code>, I noticed that the my (single) video was being loaded many times, when really I just wanted that single frame to be repeated many times.</p>
<p>Note, I’m not sure if the code accurately reproduces what I’m doing, so it might be that I am using cycle correctly in the given example but previously I was using it incorrectly.</p>

nivek: <p><code>.cycle</code> does indeed re-do all the operations before <code>n</code> times, one workaround is to put <a href="https://pytorch.org/data/beta/generated/torchdata.datapipes.iter.InMemoryCacheHolder.html#torchdata.datapipes.iter.InMemoryCacheHolder" rel="noopener nofollow ugc"><code>.in_memory_cache(size=BUFFER_SIZE)</code></a> after it. Otherwise, <code>.flatmap()</code> is likely a better option here.</p>
<p>I can see how this is not the best user experience. I am opening a GitHub issue for discussion.</p>

Vedant_Roy: <p>Does the <code>shuffle</code> (which has a buffer), also use an in-memory cache? I assume it has to.</p>

nivek: <p>Elements are popped off the buffer in <code>shuffle</code> as soon as they are yielded, whereas the elements in <code>in_memory_cache</code> persist. Therefore, using the latter should mean you will not be reading from the disk more than once unless you cache size is too small relative to the amount of data you’d like to store.</p>

Vedant_Roy: <p><a class="mention" href="/u/nivek">@nivek</a> I’m having a hard time understanding this. Presumably, for shuffling to work, it is something like</p>
<ol>
<li>Wait till the shuffle buffer is full (has e.g 10K elements)</li>
<li>When the destination datapipe needs an element, pop a random element out of the buffer, and then replace it with an element from the source datapipe</li>
<li>If the destination source datapipe has no more elements just yield elements until the buffer is empty</li>
</ol>
<p>It seems like the result of this would be that there is effectively an in-memory cache of 10K elements.</p>

Vedant_Roy: <p><a class="mention" href="/u/nivek">@nivek</a> I’ve been trying to understand the <code>cycle</code> implementation so I can figure out things like this in the future and it doesn’t quite make sense to me. Can you give me a few pointers?</p>
<pre><code class="lang-auto">    def __iter__(self) -&gt; Iterator[T_co]:
        i = 0
        while self.count is None or i &lt; self.count:
            yield from self.source_datapipe
            i += 1
</code></pre>
<p>From reading this, it looks like it will do a <code>yield from</code> the source data pipe a total of <code>count</code> times.<br>
The thing that is confusing is, (let’s say <code>count == 3</code> is 3), then I would expect <code>cycle</code> to yield from the source data pipe 3 times.</p>
<p>But, if the source data pipe has a total of 3 elements then I would expect <code>cycle</code> to yield <code>9</code> elements, not 3. Does my confusion here make sense?</p>

nivek: <p>Your description of buffer shuffle is basically accurate. As you describe in step 3, it yields element from the buffer (i.e. it pops). As soon as it yields, the element is removed from the buffer. If you want that element again, you have to re-do the previous operations, such as reading from files and etc.</p>
<p>Whereas for a cache, an element is not evicted upon reading. It stays there, such that you can read the same element multiple times without eviction as long as it has enough space. Does that make sense?</p>

nivek: <p>Edit: Sorry previous description is slightly off.</p>
<p>The line <code>yield from self.source_datapipe</code> will run 9 times, and each time it is ran it will yield an element from the DataPipe. After 3 yields, the source DataPipe is exhausted and we move onto <code>i += 1</code>. In the next iteration, it resets the source DataPipe and does the same thing. Overall, it repeats <code>n_cycle</code> times, so 3 x 3 = 9 elements. Does that answer your question?</p>

nivek: <p>You can basically re-write that code snippet as:</p>
<pre><code class="lang-python">i = 0
while self.count is None or i &lt; self.count:
    for e in self.source_datapipe:
        yield e
    i += 1
</code></pre>
<p>I think your confusion might be around what happens when you call <code>iter(datapipe)</code> multiple times. Basically every time <code>iter(datapipe)</code> is called (explicitly or implicitly), it most likely also call on its source DataPipe as well, hence operations need to be re-done. It saves memory but at the cost of needing to re-do operations. You can avoid this by using the various caching mechanisms that we provide.</p>

Vedant_Roy: <p>Thanks for the detailed explanations! The <code>yield from</code> makes sense now (I think). Basically, it ‘passes through’ the underlying generator from the source datapipe.</p>
<p>In the case of cycle, we “pass through” the underlying generator “count” times.</p>
<p><strong>Edit:</strong> nevermind, think I figured it out<br>

I’m still a bit confused about the shuffle buffer vs cache b/c it just seems like the difference is a difference in ordering. I.e., the shuffle buffer pops out elements randomly while the cache is FIFO. I’ve been looking at the cache code and it doesn’t really make sense to me. I might try &amp; re-implement it, and see if that helps.
</p>
<p>Thanks again for all the responses!</p>