pytorch dataloader中每个线程随机种子问题

今天调代码遇到一个有趣的现象,在全局设置好随机种子后,在不同的num_workers中,numpy,random,torch中生成的随机数有差别,比如下面

#num_wokers=3时候,分别打印print(random.random(),torch.rand(1),np.random.rand()),
#可以发现用np随机出来的值竟然是一样的,所以我们如果要想让np随机出来的值不一样,需要自己设置一个回调函数worker_init_fn

0.5682499026693681 tensor([0.9270]) 0.6221087710398319
0.05174456999040877 tensor([0.6482]) 0.6221087710398319
0.34046088850898915 tensor([0.1377]) 0.6221087710398319
#dp模式下种子
def worker_init_fn(worker_id):
    seed = torch.utils.data.get_worker_info().seed%(2**32-1)
    np.random.seed(seed)
#ddp 模式下的随机种子
   if args.seed is not None:
        print("Using seed = {}".format(args.seed))
        torch.manual_seed(args.seed + args.local_rank)
        torch.cuda.manual_seed(args.seed + args.local_rank)
        np.random.seed(seed=args.seed + args.local_rank)
        random.seed(args.seed + args.local_rank)
        def worker_init_fn(id):
            np.random.seed(seed=args.seed + args.local_rank + id)
            random.seed(args.seed + args.local_rank + id)

dataloader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=3, worker_init_fn=worker_init_fn)