聊聊Pytorch中的dataloader

今天為啥突然要寫一下pytorch的dataloader呢，首先來說說事情的來龍去脈。

起初，我最開始單獨訓練一個網(wǎng)絡來完成landmark點回歸任務和分類任務，訓練的數(shù)據(jù)是txt格式，在訓練之前對數(shù)據(jù)進行分析，發(fā)現(xiàn)分類任務中存在嚴重的數(shù)據(jù)樣本不均衡的問題，那么我事先針對性的進行數(shù)據(jù)采樣均衡操作，重新得到訓練和測試的txt數(shù)據(jù)和標簽，保證了整個訓練和測試數(shù)據(jù)的樣本均衡性。由于我的整個項目是檢測+點回歸+分類，起初檢測和點回歸+分類是分兩步實現(xiàn)的，檢測是通過讀取XML格式來進行訓練，現(xiàn)在要統(tǒng)一整個項目的訓練和測試過程，要將點回歸+分類的訓練測試過程也按照讀取XML格式來進行，那么就遇到一個問題，如何針對性的去給樣本偏少的樣本進行均衡，由于在dataset類中，返回的圖像和標簽都是針對每個index返回一個結(jié)果，在dataset類中進行操作似乎不太可行，那么就想到在dataloader中進行操作，通過dataloader中的參數(shù)sample來完成針對性采樣。

還有一個問題是關于num_workers的設置，因為我有對比過，在我的單機RTX 2080Ti上和八卡服務器TITAN RTX上(僅使用單卡，其它卡有在跑其它任務)，使用相同的num_workers，在單機上的訓練速度反而更快，于是猜想可能和CPU或者內(nèi)存有關系，下面會具體分析。

首先來看下下dataloader中的各個參數(shù)的含義。

類的定義為：torch.utils.data.DataLoader ，其中包含的參數(shù)有：

torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, \    batch_sampler=None, num_workers=0, collate_fn=None, pin_memory=False, \    drop_last=False, timeout=0, worker_init_fn=None, multiprocessing_context=None)

dataset：定義的dataset類返回的結(jié)果。

batchsize：每個bacth要加載的樣本數(shù)，默認為1。

shuffle：在每個epoch中對整個數(shù)據(jù)集data進行shuffle重排，默認為False。

sample：定義從數(shù)據(jù)集中加載數(shù)據(jù)所采用的策略，如果指定的話，shuffle必須為False；batch_sample類似，表示一次返回一個batch的index。

num_workers：表示開啟多少個線程數(shù)去加載你的數(shù)據(jù)，默認為0，代表只使用主進程。

collate_fn：表示合并樣本列表以形成小批量的Tensor對象。

pin_memory：表示要將load進來的數(shù)據(jù)是否要拷貝到pin_memory區(qū)中，其表示生成的Tensor數(shù)據(jù)是屬于內(nèi)存中的鎖頁內(nèi)存區(qū)，這樣將Tensor數(shù)據(jù)轉(zhuǎn)義到GPU中速度就會快一些，默認為False。

drop_last：當你的整個數(shù)據(jù)長度不能夠整除你的batchsize，選擇是否要丟棄最后一個不完整的batch，默認為False。

注：這里簡單科普下pin_memory，通常情況下，數(shù)據(jù)在內(nèi)存中要么以鎖頁的方式存在，要么保存在虛擬內(nèi)存(磁盤)中，設置為True后，數(shù)據(jù)直接保存在鎖頁內(nèi)存中，后續(xù)直接傳入cuda；否則需要先從虛擬內(nèi)存中傳入鎖頁內(nèi)存中，再傳入cuda，這樣就比較耗時了，但是對于內(nèi)存的大小要求比較高。

下面針對num_workers，sample和collate_fn分別進行說明：

1. 設置num_workers：

pytorch中dataloader一次性創(chuàng)建num_workers個子線程，然后用batch_sampler將指定batch分配給指定worker，worker將它負責的batch加載進RAM，dataloader就可以直接從RAM中找本輪迭代要用的batch。如果num_worker設置得大，好處是尋batch速度快，因為下一輪迭代的batch很可能在上一輪/上上一輪...迭代時已經(jīng)加載好了。壞處是內(nèi)存開銷大，也加重了CPU負擔（worker加載數(shù)據(jù)到RAM的進程是進行CPU復制）。如果num_worker設為0，意味著每一輪迭代時，dataloader不再有自主加載數(shù)據(jù)到RAM這一步驟，只有當你需要的時候再加載相應的batch，當然速度就更慢。num_workers的經(jīng)驗設置值是自己電腦/服務器的CPU核心數(shù)，如果CPU很強、RAM也很充足，就可以設置得更大些，對于單機來說，單跑一個任務的話，直接設置為CPU的核心數(shù)最好。

2. 定義sample：（假設dataset類返回的是：data, label）

from torch.utils.data.sampler import WeightedRandomSampler## 如果label為1，那么對應的該類別被取出來的概率是另外一個類別的2倍weights = [2 if label == 1 else 1 for data, label in dataset]sampler = WeightedRandomSampler(weights,num_samples=10, replacement=True)dataloader = DataLoader(dataset, batch_size=16, sampler=sampler)

PyTorch中提供的這個sampler模塊，用來對數(shù)據(jù)進行采樣。默認采用SequentialSampler，它會按順序一個一個進行采樣。常用的有隨機采樣器：RandomSampler，當dataloader的shuffle參數(shù)為True時，系統(tǒng)會自動調(diào)用這個采樣器，實現(xiàn)打亂數(shù)據(jù)。這里使用另外一個很有用的采樣方法：WeightedRandomSampler，它會根據(jù)每個樣本的權重選取數(shù)據(jù)，在樣本比例不均衡的問題中，可用它來進行重采樣。replacement用于指定是否可以重復選取某一個樣本，默認為True，即允許在一個epoch中重復采樣某一個數(shù)據(jù)。

3. 定義collate_fn：

def detection_collate(batch):    """Custom collate fn for dealing with batches of images that have a different    number of associated object annotations (bounding boxes).
    Arguments:        batch: (tuple) A tuple of tensor images and lists of annotations
    Return:        A tuple containing:            1) (tensor) batch of images stacked on their 0 dim            2) (list of tensors) annotations for a given image are stacked on                                 0 dim    """    targets = []    imgs = []    for sample in batch:        imgs.append(sample[0])        targets.append(torch.FloatTensor(sample[1]))    return torch.stack(imgs, 0), targets

使用dataloader時加入collate_fn參數(shù)，即可合并樣本列表以形成小批量的Tensor對象，如果你的標簽不止一個的話，還可以支持自定義，在上述方法中再額外添加對應的label即可。

data_loader = torch.utils.data.DataLoader(dataset, args.batch_size,    num_workers=args.num_workers, sampler=sampler, shuffle=False,     collate_fn=detection_collate, pin_memory=True, drop_last=True)

參考鏈接：