pytorch中關於distributedsampler函數的使用
關於distributedsampler函數的使用
1.如何使用這個分佈式采樣器
在使用distributedsampler函數時,觀察loss發現loss收斂有規律,發現是按順序讀取數據,未進行shuffle。
問題的解決方式就是懷疑 seed 有問題,參考源碼 DistributedSampler,發現 shuffle 的結果依賴 g.manual_seed(self.epoch) 中的 self.epoch。
def __iter__(self): # deterministically shuffle based on epoch g = torch.Generator() g.manual_seed(self.epoch) if self.shuffle: indices = torch.randperm(len(self.dataset), generator=g).tolist() else: indices = list(range(len(self.dataset))) # add extra samples to make it evenly divisible indices += indices[:(self.total_size - len(indices))] 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)
而 self.epoch 初始默認是 0
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 self.shuffle = shuffle
但是 DistributedSampler 也提供瞭一個 set 函數來改變 self.epoch
def set_epoch(self, epoch): self.epoch = epoch
所以在運行的時候要不斷調用這個 set_epoch 函數。隻要把我的代碼中的
# sampler.set_epoch(e)
全部代碼如下:
import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler torch.distributed.init_process_group(backend="nccl") input_size = 5 output_size = 2 batch_size = 2 data_size = 16 local_rank = torch.distributed.get_rank() torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) class RandomDataset(Dataset): def __init__(self, size, length, local_rank): self.len = length self.data = torch.stack([torch.ones(5), torch.ones(5)*2, torch.ones(5)*3,torch.ones(5)*4, torch.ones(5)*5,torch.ones(5)*6, torch.ones(5)*7,torch.ones(5)*8, torch.ones(5)*9, torch.ones(5)*10, torch.ones(5)*11,torch.ones(5)*12, torch.ones(5)*13,torch.ones(5)*14, torch.ones(5)*15,torch.ones(5)*16]).to('cuda') self.local_rank = local_rank def __getitem__(self, index): return self.data[index] def __len__(self): return self.len dataset = RandomDataset(input_size, data_size, local_rank) sampler = DistributedSampler(dataset) rand_loader = DataLoader(dataset=dataset, batch_size=batch_size, sampler=sampler) e = 0 while e < 2: t = 0 # sampler.set_epoch(e) for data in rand_loader: print(data) e+=1
運行:
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 test.py
2.關於用不用這個采樣器的區別
多卡去訓模型,嘗試著用DDP模式,而不是DP模式去加速訓練(很容易出現負載不均衡的情況)。
遇到瞭一點關於DistributedSampler這個采樣器的一點疑惑,想試驗下在DDP模式下,使用這個采樣器和不使用這個采樣器有什麼區別。
實驗代碼:
整個數據集大小為8,batch_size 為4,總共跑2個epoch。
import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler torch.distributed.init_process_group(backend="nccl") batch_size = 4 data_size = 8 local_rank = torch.distributed.get_rank() print(local_rank) torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) class RandomDataset(Dataset): def __init__(self, length, local_rank): self.len = length self.data = torch.stack([torch.ones(1), torch.ones(1)*2,torch.ones(1)*3,torch.ones(1)*4,torch.ones(1)*5,torch.ones(1)*6,torch.ones(1)*7,torch.ones(1)*8]).to('cuda') self.local_rank = local_rank def __getitem__(self, index): return self.data[index] def __len__(self): return self.len dataset = RandomDataset(data_size, local_rank) sampler = DistributedSampler(dataset) #rand_loader =DataLoader(dataset=dataset,batch_size=batch_size,sampler=None,shuffle=True) rand_loader = DataLoader(dataset=dataset,batch_size=batch_size,sampler=sampler) epoch = 0 while epoch < 2: sampler.set_epoch(epoch) for data in rand_loader: print(data) epoch+=1
運行命令:
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 test.py
實驗結果:
結論分析:上面的運行結果來看,在一個epoch中,sampler相當於把整個數據集 劃分成瞭nproc_per_node份,每個GPU每次得到batch_size的數量,也就是nproc_per_node 個GPU分一整份數據集,總數據量大小就為1個dataset。
如果不用它裡面自帶的sampler,單純的還是按照我們一般的形式。Sampler=None,shuffle=True這種,那麼結果將會是下面這樣的:
結果分析:沒用sampler的話,在一個epoch中,每個GPU各自維護著一份數據,每個GPU每次得到的batch_size的數據,總的數據量為2個dataset,
總結
一般的形式的dataset隻能在同進程中進行采樣分發,也就是為什麼圖2隻能單GPU維護自己的dataset,DDP中的sampler可以對不同進程進行分發數據,圖1,可以誇不同進程(GPU)進行分發。
以上為個人經驗,希望能給大傢一個參考,也希望大傢多多支持WalkonNet。
推薦閱讀:
- Pytorch distributed 多卡並行載入模型操作
- 詳解pytorch的多GPU訓練的兩種方式
- pytorch DistributedDataParallel 多卡訓練結果變差的解決方案
- pytorch DataLoader的num_workers參數與設置大小詳解
- 我對PyTorch dataloader裡的shuffle=True的理解