Checkpoints not getting created with custom sampler

username: Rini
cooked: <p>Hi,</p>
<p>I am working on a multi-task model with uneven dataset size and have a custom sampler and using the sampler in dataloader (below)</p>
<pre><code>sampler = BalancedBatchSchedulerSampler(dataset, batch_size)
dataloader = DataLoader(
                dataset,
                sampler=sampler,
                batch_size=batch_size,
                collate_fn=collate_fn,
                num_workers=num_workers,
            )
</code></pre>
<p><code>BalancedBatchSchedulerSampler</code> is the custom sampler. Also I have set <code>replace_sampler_ddp</code> to <code>False</code>. With this custom sampler I don’t see the checkpoint folder is getting created for the model. When I don’t pass the <code>sampler</code> argument and use the default <code>RandomSampler</code> the checkpoint is getting created without any other change in the code.</p>
<p>Is it possible that the sampler is affecting the model checkpoint somehow?</p>
<p>Thank you!</p>

username: mMagmer
cooked: <p>Maybe your custom sampler somehow has very large size, And you won’t reach the end of it. check the implemention of <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.Sampler" rel="noopener nofollow ugc"><code>__len__</code></a> method.</p>

username: Rini
cooked: <aside class="quote no-group" data-username="mMagmer" data-post="2" data-topic="139699">
<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/mmagmer/48/43283_2.png" class="avatar"> mMagmer:</div>
<blockquote>
<p><strong>len</strong></p>
</blockquote>
</aside>
<p>Thanks so much for pointing to the <code>__len__</code> method. I have a multi-task model and the three tasks have different sizes. The <code>__len__</code> method is implemented to return equal proportion of samples from each dataset. I went through PyTorch documentation on Sampler and dataloader but not sure what should I change. Did you mean that the number of samples are too large?</p>
<p>Thank you</p>
<pre><code>import math
from random import shuffle
import torch
from torch.utils.data.sampler import RandomSampler


class BalancedBatchSchedulerSampler(torch.utils.data.sampler.Sampler):
    """
    iterate over tasks and provide a balanced batch per task in each mini-batch
    """

    def __init__(self, dataset, batch_size):

        super(BalancedBatchSchedulerSampler, self).__init__(dataset)
        self.dataset = dataset
        self.batch_size = batch_size
        self.number_of_datasets = len(dataset.datasets)
        self.largest_dataset_size = max(
            [len(cur_dataset) for cur_dataset in dataset.datasets]
        )

    def __len__(self):
        return (
            self.batch_size
            * math.ceil(self.largest_dataset_size / self.batch_size)
            * len(self.dataset.datasets)
        )

    def __iter__(self):
        samplers_list = []
        sampler_iterators = []
        for dataset_idx in range(self.number_of_datasets):
            cur_dataset = self.dataset.datasets[dataset_idx]
            sampler = RandomSampler(cur_dataset)
            samplers_list.append(sampler)
            cur_sampler_iterator = sampler.__iter__()
            sampler_iterators.append(cur_sampler_iterator)

        push_index_val = [0] + list(self.dataset.cumulative_sizes[:-1])
        step = self.batch_size
        samples_to_grab = math.ceil(self.batch_size / self.number_of_datasets)
        # for this case we want to get all samples in dataset, this force us to resample from the smaller datasets
        epoch_samples = self.largest_dataset_size * self.number_of_datasets

        final_samples_list = []  # this is a list of indexes from the combined dataset
        for _ in range(0, epoch_samples, step):
            cur_batch_samples = []
            for i in range(self.number_of_datasets):
                cur_batch_sampler = sampler_iterators[i]
                cur_samples = []
                for _ in range(samples_to_grab):
                    try:
                        cur_sample_org = cur_batch_sampler.__next__()
                        cur_sample = cur_sample_org + push_index_val[i]
                        cur_samples.append(cur_sample)
                    except StopIteration:
                        # got to the end of iterator - restart the iterator and continue to get samples
                        # until reaching "epoch_samples"
                        sampler_iterators[i] = samplers_list[i].__iter__()
                        cur_batch_sampler = sampler_iterators[i]
                        cur_sample_org = cur_batch_sampler.__next__()
                        cur_sample = cur_sample_org + push_index_val[i]
                        cur_samples.append(cur_sample)
                cur_batch_samples.extend(cur_samples)
            shuffle(cur_batch_samples)
            final_samples_list.extend(cur_batch_samples)

        return iter(final_samples_list)
</code></pre>

username: mMagmer
cooked: <aside class="quote no-group" data-username="Rini" data-post="3" data-topic="139699">
<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/rini/48/32526_2.png" class="avatar"> Rini:</div>
<blockquote>
<p>Did you mean that the number of samples are too large?</p>
</blockquote>
</aside>
<p>That was my guess.<br>
Based on the code, I think you’re implementing a BatchSampler not a sampler.<br>
I don’t know how you can fix your code,</p>

username: Rini
cooked: <p>I turned off the <code>__len__</code> function and checkpoint is getting created. I think I will play around a bit with batch sampler.<br>
Thanks again for the help!</p>

username: mMagmer
cooked: <p>maybe you can get what you want by playing with batchsize,<br>
something like:</p>
<pre><code class="lang-auto">import torch
d1 = torch.utils.data.TensorDataset(torch.ones(100,2))
d2 = torch.utils.data.TensorDataset(2*torch.ones(200,3))
d3 = torch.utils.data.TensorDataset(3*torch.ones(300,4))
dl1 = torch.utils.data.DataLoader(dataset=d1, batch_size=5, shuffle=True)
dl2 = torch.utils.data.DataLoader(dataset=d2, batch_size=10, shuffle=True)
dl3 = torch.utils.data.DataLoader(dataset=d3, batch_size=15, shuffle=True)
dl =zip(dl1,dl2,dl3)
#next(iter(dl))
#len(list(dl))
for m,n,p in dl:
    # do the job
</code></pre>