Pytorch中DistributedSampler()中的随机因素

在使用YOLOv5的mosaic数据增强（从数据集中随机采样3张图像与当前的拼在一起，组成一张里面有4张小图的大图）过程中，当使用多卡（DDP）训练自己的算法时出现了每张卡所生成的图像中，除了原本的那张小图，其他三张小图都一模一样的现象：

可以看到除了左上角的小图以外，其余三张采样的图片完全相同。这样的同态性的数据降低了训练数据的质量，我们希望用于数据增强的三张小图也是不同的。查看YOLOv5源码中mosaic数据增强这一块源码：

class LoadImagesAndLabels(Dataset):  # for training/testing
    def __init__()
    ...
    def __getitem__(self, index):
        index = self.indices[index]  # linear, shuffled, or image_weights

        hyp = self.hyp
        mosaic = self.mosaic and random.random() < hyp['mosaic']
        if mosaic:
            # Load mosaic
            img, labels = load_mosaic(self, index)
            ...

def load_mosaic(self, index):
    # loads images in a mosaic

    labels4 = []
    s = self.img_size
    yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border]  # mosaic center x, y
    indices = [index] + [self.indices[random.randint(0, self.n - 1)] for _ in range(3)]  # 3 additional image indices
    ...

可以看到源码中这块采用了random来进行随机抽取其余三张小图，那是不是我们在每个子进程中为random固定了相同的随机数种子的原因，所以我们为每张卡固定不同的随机数种子：

def main_worker(gpu, ngpus_per_node, args):

    #torch.multiprocessing.set_sharing_strategy('file_system')
    random.seed(args.seed+gpu)#to use yolov5-mosaic
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)  
    torch.cuda.manual_seed(args.seed)  
    torch.cuda.manual_seed_all(args.seed)
    ...

这样每张卡上random随机数种子就是不一样的，按理说采样出来的三张小图就是不一样的了，但是结果跟刚刚还是一样：除了本身的那张小图，其余三张小图还是一样。。。经过反复查找问题，定位在torch.utils.data.distributed.DistributedSampler()上，先来看看pytorch中如何解释DistributedSampler

看看其中几个关键的参数：

• dataset 不用多说
• num_replicas 参与分布式训练的进程数，默认情况下将自己检索这个参数
• rank 当前进程的编号，默认情况下将自己检索这个参数
• shuffle 是否打乱数据集，默认为True
• seed 当shuffle=True时，用于random的随机数种子。这个参数在所有分布式进程中必须保持一致，默认为0

划重点！！！应该就是DistributedSampler中的seed参数，它会在取数据时将random的随机数种子全部默认设置为0，这样每张卡采样的其余三张小图当然就一模一样了。那为什么在DistributedSampler中要将random.seed固定为相同的呢，放上我的理解:

比方说我们用四卡训练，总的bs=16,那么平均到每张卡上就是放4张图片。DistributedSampler默认shuffle=True,那么从哪个位置开始取16张图片呢，这就要由random来随机确定，如果所有进程的随机数种子是一样的，那么所有进程每次确定的这个起点就是相同的，比方说当前iter所有卡都从1000开始取数据，rank=0的进程取[1000,1001,1002,1003],rank=1的进程取[1004,1005,1006,1007],rank=2取[1008,1009,1010,1011]，rank=3取[1012,1013,1014,1015]。如果seed不相同，那么每个iter就可能出现rank=0的进程从1000开始取数据，rank=1的进程从2000开始取数据的这种情况，这当然是不合理的。

sampler = DistributedSampler(dataset) if is_distributed else None
loader = DataLoader(dataset, shuffle=(sampler is None),sampler=sampler)
for epoch in range(start_epoch, n_epochs):
    if is_distributed:
	sampler.set_epoch(epoch)
	train(loader)

同时由官网的这个例子可以看到，因为在用DistributedSampler实例化sampler的时候就已经shuffle过数据了，所以在定义dataloader的时候不需要再将shuffle设置为True

好了，回归mosaic，如何让每个iter中所有进程取出的其余三张小图不同呢，用np.random不就解决了！！在dataset定义过程中将random生成的随机数换成用np.random生成,再在每个进程中为np.random.seed赋予不同的随机数种子，这样每个进程每次在数据处理上的随机因素就不同了。代码修改如下：

def load_mosaic(self, index):
    # loads images in a mosaic

    labels4 = []
    s = self.img_size
    yc, xc = [int(np.random.uniform(-x, 2 * s + x)) for x in self.mosaic_border]  # mosaic center x, y
    indices = [index] + [self.indices[np.random.randint(0, self.n)] for _ in range(3)]  # 3 additional image indices
    ...

def main_worker(gpu, ngpus_per_node, args):

    #torch.multiprocessing.set_sharing_strategy('file_system')
    random.seed(args.seed)
    np.random.seed(args.seed+gpu) # for yolov5-mosaic
    torch.manual_seed(args.seed)  
    torch.cuda.manual_seed(args.seed) 
    torch.cuda.manual_seed_all(args.seed)

最后再看看一个iter下每张卡取出的数据：

ok，问题解决。