PyTorch 分布式训练：详解自定义 Sampler 类代码

在 PyTorch 分布式训练中，为了高效地在多个进程间分配和采样数据，我们需要使用 torch.utils.data.Sampler 类或其子类。以下代码定义了一个自定义的 Sampler 类，用于在分布式环境下进行数据采样：pythonimport torchimport numpy as np

class Sampler(torch.utils.data.Sampler): def init(self, dataset, num_replicas=None, rank=None, shuffle=True, make_even=True): if num_replicas is None: if not torch.distributed.is_available(): raise RuntimeError('Requires distributed package to be available') num_replicas = torch.distributed.get_world_size() if rank is None: if not torch.distributed.is_available(): raise RuntimeError('Requires distributed package to be available') rank = torch.distributed.get_rank() self.shuffle = shuffle self.make_even = make_even self.dataset = dataset self.num_replicas = num_replicas self.rank = rank self.epoch = 0 self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas)) self.total_size = self.num_samples * self.num_replicas indices = list(range(len(self.dataset))) self.valid_length = len(indices[self.rank : self.total_size : self.num_replicas])

def __iter__(self):        if self.shuffle:            g = torch.Generator()            g.manual_seed(self.epoch)            indices = torch.randperm(len(self.dataset), generator=g).tolist()        else:            indices = list(range(len(self.dataset)))        if self.make_even:            if len(indices) < self.total_size:                if self.total_size - len(indices) < len(indices):                    indices += indices[: (self.total_size - len(indices))]                else:                    extra_ids = np.random.randint(low=0, high=len(indices), size=self.total_size - len(indices))                    indices += [indices[ids] for ids in extra_ids]            assert len(indices) == self.total_size        indices = indices[self.rank : self.total_size : self.num_replicas]        self.num_samples = len(indices)        return iter(indices)

def __len__(self):        return self.num_samples

def set_epoch(self, epoch):        self.epoch = epoch

代码解析:

初始化方法 (__init__): - 首先检查是否支持分布式训练，并获取进程数 (num_replicas) 和当前进程的 rank (rank)。 - 设置是否打乱数据 (shuffle) 和使数据集长度均匀 (make_even) 的参数。 - 保存数据集、进程数、当前进程的 rank、当前 epoch 等信息。 - 计算每个进程需要采样的样本数 (num_samples) 和总的样本数 (total_size)。 - 根据 rank 和进程数计算当前进程需要处理的样本索引范围 (valid_length)。
迭代器方法 (__iter__): - 根据 shuffle 参数决定是否打乱数据集索引。 - 根据 make_even 参数决定是否对数据集长度进行调整，使其可以被进程数整除。 - 根据 rank 和进程数获取当前进程需要处理的样本索引。 - 返回一个迭代器，用于遍历当前进程的样本索引。
长度方法 (__len__): - 返回当前进程需要处理的样本数。
设置 epoch 方法 (set_epoch): - 用于设置当前 epoch 的值，主要用于在每个 epoch 开始时更新随机数生成器的种子，确保每个 epoch 的数据采样不同。

总结:

这个自定义的 Sampler 类可以帮助我们在 PyTorch 分布式训练中进行高效的数据采样。通过设置不同的参数，我们可以控制是否打乱数据、是否使数据集长度均匀等，从而灵活地满足不同的训练需求。