How to create a DistributedSampler based on my own Sampler

username: Boltzmachine
cooked: <p>I have implement a <code>Sampler</code> when using single GPU. I implement it because I want all samples in a batch are from the same source.</p>
<pre><code class="lang-auto">class MultiSourceBatchSampler(torch.utils.data.Sampler):
    def __init__(
        self,
        data_source: Dataset,
        # sources: list,
        batch_size: int,
        shuffle: bool = False,
        drop_last: bool = False,
    ):
        self.data_source = data_source
        # self.sources = sources
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.drop_last = drop_last

        self._length = -1
    
    def __iter__(self):
        # indices are tuples (source, second_key)
        indices = self.data_source._get_indices()
        if self.shuffle:
            np.random.shuffle(indices)
        indices.sort(key=lambda x: x[0]) #sort by source

        batched_indices = []
        buffer = []
        for index in indices:
            if (len(buffer) == self.batch_size) or (len(buffer) &gt; 0 and buffer[0][0] != index[0]):
                batched_indices.append(buffer)
                buffer = []
            buffer.append(index)
        if len(buffer) &gt; 0 and not self.drop_last:
            batched_indices.append(buffer)
       
        if self.shuffle:
            np.random.shuffle(batched_indices)
        yield from batched_indices

    def __len__(self):
        if self._length &gt;= 0:
            return self._length
        total = 0
        indices = self.data_source._get_keys(by=self.by)
        indices.sort(key=lambda x: x[0]) #sort by source

        buffer = []
        for index in indices:
            if (len(buffer) == self.batch_size) or (len(buffer) &gt; 0 and buffer[0][0] != index[0]):
                total += 1
                buffer = []
            buffer.append(index)
        if len(buffer) &gt; 0 and not self.drop_last:
            total += 1
        self._length = total
        return total
</code></pre>
<p>I did not find very detailed documentation about <code>DistributedSampler</code>, so I wonder what is the principle if I want to convert my sampler to a DistributedSampler</p>

username: ptrblck
cooked: <p>The <a href="https://github.com/pytorch/pytorch/blob/c7dbee51066e27e0a13c3877de332002fde9e125/torch/utils/data/distributed.py#L65-L133"><code>DistributedSampler</code></a> splits the dataset indices based on the number of replicas making sure each rank receives the same number of samples. The actual <code>self.total_size</code> depends on the number of samples in the dataset, the number of replicas and if the last samples should be dropped or repeated as seen <a href="https://github.com/pytorch/pytorch/blob/c7dbee51066e27e0a13c3877de332002fde9e125/torch/utils/data/distributed.py#L115-L126">here</a>. It also holds the <code>seed</code> allowing you to reset the <code>epoch</code> to continue training from the same epoch.</p>

username: Boltzmachine
cooked: <p>A possible general solution:</p>
<pre data-code-wrap="python"><code class="lang-python">class DistributedSamplerWrapper(torch.utils.data.distributed.DistributedSampler):
    def __init__(
        self,
        base_sampler,
        num_replicas = None,
        rank = None,
        seed = 0,
    ):
        self.base_sampler = base_sampler
        self.batch_size = base_sampler.batch_size
        shuffle = self.base_sampler.shuffle
        drop_last = self.base_sampler.drop_last
        super().__init__(base_sampler, num_replicas=num_replicas, rank=rank, shuffle=shuffle, seed=seed, drop_last=drop_last)

    def __iter__(self):
        base_indices = list(self.base_sampler.__iter__())
        if self.shuffle:
            # deterministically shuffle based on epoch and seed
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)
            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
        else:
            indices = list(range(len(self.dataset)))  # type: ignore[arg-type]
        indices = [base_indices[i] for i in indices]
            
        if not self.drop_last:
            # add extra samples to make it evenly divisible
            padding_size = self.total_size - len(indices)
            if padding_size &lt;= len(indices):
                indices += indices[:padding_size]
            else:
                indices += (indices * math.ceil(padding_size / len(indices)))[
                    :padding_size
                ]
        else:
            # remove tail of data to make it evenly divisible.
            indices = indices[: self.total_size]
        assert len(indices) == self.total_size

        # subsample
        indices = indices[self.rank : self.total_size : self.num_replicas]
        assert len(indices) == self.num_samples

        return iter(indices)
</code></pre>
<p>Should be passed as the <code>batch_sampler</code> argument.</p>