Is there any way to skip steps in a DataLoader?

username: Alexander_Soare
cooked: <p>A couple of times I’ve been training on large datasets where one epoch takes several hours and a 1 in 1,000,000 bug came up at some point during the epoch. Luckily I validate every so number of train steps and at this point I save a checkpoint.</p>
<p>BUT, the problem is that I want to get back to exactly where I was previously. So in my train loop I can say</p>
<pre><code class="lang-python">if train_step &lt;= n:
    continue
</code></pre>
<p>But then my dataloader is still doing all the I/O and preprocessing for all steps &lt;= n. I then tried this in my dataset</p>
<pre><code class="lang-python">class Dataset(torch.utils.data.Dataset):
    def __init__(self):
        ....
        self.skip = True

    def __getitem__(self):
        if self.skip:
            return torch.zeros(1)
        # normal I/O and preprocessing
</code></pre>
<p>Then in my train loop:</p>
<pre><code class="lang-auto">if train_step &lt;= n:
    if train_step == n:
        dataset.skip = False
    continue
</code></pre>
<p>But this seems to have no effect.</p>
<p><strong>Is there a way to set a certain step index on a DataLoader prior to iteration?</strong></p>

username: ejguan
cooked: <p>Maybe try to use custom Sampler?</p>

username: Alexander_Soare
cooked: <p>I suppose so. There are many things I could do. Was just wondering if there was a built-in way to do this easily as I imagine it’s not an uncommon need.</p>

username: ejguan
cooked: <p>Yeah, I would say no built-in way for now. But, we are working on a new design of DataLoader, which IMO will provide this functionality.</p>

username: doem97
cooked: <p>Hi <a class="mention" href="/u/alexander_soare">@Alexander_Soare</a> so finally how do you make it? And work around maybe helpful!<br>
Thanks.</p>

username: Alexander_Soare
cooked: <p><a class="mention" href="/u/doem97">@doem97</a> sorry but no I didn’t manage to get a good workaround for this. Just ended up dodging the problem at a higher level outside the scope of this thread.</p>

username: satpalsingh
cooked: <p>Any latest solution for this? <a class="mention" href="/u/ejguan">@ejguan</a></p>

username: ejguan
cooked: <aside class="quote no-group" data-username="Alexander_Soare" data-post="1" data-topic="123201">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alexander_soare/48/24720_2.png" class="avatar"> Alexander_Soare:</div>
<blockquote>
<p>ain step index on a DataLoa</p>
</blockquote>
</aside>
<p>The feature of <code>DataLoader2</code> to do checkpoint over data-pipeline is still WIP.</p>

username: ejguan
cooked: <p>As a workaround for people requesting it, I think you can do it via a custom sampler</p>
<pre data-code-wrap="py"><code class="lang-nohighlight">class CustomSampler(Sampler[int]):
    def __init__(self, data_source):
        self.data_source = data_source
        self.skip_ctr = 0

    def skip(self, n):
        self.skip_ctr = n

    def __iter__(self):
        n = len(self.data_source)
        for i in range(self.skip_ctr, n):
            yield i

    def __len__(self):
        return n - self.skip_ctr
</code></pre>