Pytorch--torch.utils.data.DataLoader解讀


    torch.utils.data.DataLoader是Pytorch中數據讀取的一個重要接口,其在dataloader.py中定義,基本上只要是用oytorch來訓練模型基本都會用到該接口,該接口主要用來將自定義的數據讀取接口的輸出或者PyTorch已有的數據讀取接口的輸入按照batch size封裝成Tensor,后續只需要再包裝成Variable即可作為模型的輸入,主要包括DataLoader和DataLoaderIter兩個類。

    dataloader.py腳本的的github地址:https://github.com/pytorch/pytorch/blob/master/torch/utils/data/dataloader.py

    DataLoader類源碼如下。先看看__init__中的幾個重要的輸入:1、dataset,這個就是PyTorch已有的數據讀取接口(比如torchvision.datasets.ImageFolder)或者自定義的數據接口的輸出,該輸出要么是torch.utils.data.Dataset類的對象,要么是繼承自torch.utils.data.Dataset類的自定義類的對象。2、batch_size,根據具體情況設置即可。3、shuffle,一般在訓練數據中會采用。4、collate_fn,是用來處理不同情況下的輸入dataset的封裝,一般采用默認即可,除非你自定義的數據讀取輸出非常少見。5、batch_sampler,從注釋可以看出,其和batch_size、shuffle等參數是互斥的,一般采用默認。6、sampler,從代碼可以看出,其和shuffle是互斥的,一般默認即可。7、num_workers,從注釋可以看出這個參數必須大於等於0,0的話表示數據導入在主進程中進行,其他大於0的數表示通過多個進程來導入數據,可以加快數據導入速度。8、pin_memory,注釋寫得很清楚了: pin_memory (bool, optional): If True, the data loader will copy tensors into CUDA pinned memory before returning them. 也就是一個數據拷貝的問題。9、timeout,是用來設置數據讀取的超時時間的,但超過這個時間還沒讀取到數據的話就會報錯。

  在__init__中,RandomSampler類表示隨機采樣且不重復,所以起到的就是shuffle的作用。BatchSampler類則是把batch size個RandomSampler類對象封裝成一個,這樣就實現了隨機選取一個batch的目的。這兩個采樣類都是定義在sampler.py腳本中,地址:https://github.com/pytorch/pytorch/blob/master/torch/utils/data/sampler.py。以上這些都是初始化的時候進行的。當代碼運行到要從torch.utils.data.DataLoader類生成的對象中取數據的時候,比如:
train_data=torch.utils.data.DataLoader(...)
for i, (input, target) in enumerate(train_data):
...
就會調用DataLoader類的__iter__方法,__iter__方法就一行代碼:return DataLoaderIter(self),輸入正是DataLoader類的屬性。因此當調用__iter__方法的時候就牽扯到另外一個類:DataLoaderIter

  DataLoaderIter類源碼如下。self.index_queue = multiprocessing.SimpleQueue()中的multiprocessing是Python中的多進程管理包,而threading則是Python中的多線程管理包,二者很大一部分的接口用法類似。還是照例先看看__init__,前面部分都是一些賦值操作,比較特殊的是self.sample_iter = iter(self.batch_sampler),得到的self.sample_iter可以通過next(self.sample_iter)來獲取batch size個數據的index。self.rcvd_idx表示讀取到的一個batch數據的index,初始化為0,該值在迭代讀取數據的時候會用到。if self.num_workers語句是針對多進程或單進程的情況進行初始化,如果不是設置為多進程讀取數據,那么就不需要這些初始化操作,后面會介紹單進程數據讀取。在if語句中通過multiprocessing.SimpleQueue()類創建了一個簡單的隊列對象。multiprocessing.Process類就是構造進程的類,這里根據設定的進程數來啟動,然后賦值給self.workers。接下來的一個for循環就通過調用start方法依次啟動self.workers中的進程。接下來關於self.pin_memory的判斷語句,該判斷語句內部主要是實現了多線程操作。self.pin_memory的含義在前面已經介紹過了,當為True的時候,就會把數據拷到CUDA中。self.data_queue = queue.Queue()是通過Python的queue模塊初始化得到一個先進先出的隊列(queue模塊也可以初始化得到先進后出的隊列,需要用queue.LifoQueue()初始化),queue模塊主要應用在多線程讀取數據中。在threading.Thread的args參數中,第一個參數in_data就是一個進程的數據,一個進程中不同線程的數據也是通過隊列來維護的,這里采用的是Python的queue模塊來初始化得到一個隊列:queue.Queue()。初始化結束后,就會調用__next__方法,接下來介紹。
總的來說,如果設置為多進程讀取數據,那么就會采用隊列的方式來讀,如果不是采用多進程來讀取數據,那就采用普通方式來讀。

  DataLoaderIter類的__next__方法如下,包含3個if語句和1個while語句。
第一個if語句是用來處理self.num_workers等於0的情況,也就是不采用多進程進行數據讀取,可以看出在這個if語句中先通過indices = next(self.sample_iter)獲取長度為batch size的列表:indices,這個列表的每個值表示一個batch中每個數據的index,每執行一次next操作都會讀取一批長度為batch size的indices列表。然后通過self.collate_fn函數將batch size個tuple(每個tuple長度為2,其中第一個值是數據,Tensor類型,第二個值是標簽,int類型)封裝成一個list,這個list長度為2,兩個值都是Tensor,一個是batch size個數據組成的FloatTensor,另一個是batch size個標簽組成的LongTensor。所以簡單講self.collate_fn函數就是將batch size個分散的Tensor封裝成一個Tensor。batch = pin_memory_batch(batch)中pin_memory_batch函數的作用就是將輸入batch的每個Tensor都拷貝到CUDA中,該函數后面會詳細介紹。
第二個if語句是判斷當前想要讀取的batch的index(self.rcvd_idx)是否之前已經讀出來過(已讀出來的index和batch數據保存在self.reorder_dict字典中,可以結合最后的while語句一起看,因為self.reorder_dict字典的更新是在最后的while語句中),如果之前已經讀取過了,就根據這個index從reorder_dict字典中彈出對應的數據。最后返回batch數據的時候是 return self._process_next_batch(batch),該方法后面會詳細介紹。主要做是獲取下一個batch的數據index信息。
第三個if語句,self.batches_outstanding的值在前面初始中調用self._put_indices()方法時修改了,所以假設你的進程數self.num_workers設置為3,那么這里self.batches_outstanding就是3*2=6,可具體看self._put_indices()方法。
最后的while循環就是真正用來從隊列中讀取數據的操作,最主要的就是idx, batch = self._get_batch(),通過調用_get_batch()方法來讀取,后面有介紹,簡單講就是調用了隊列的get方法得到下一個batch的數據,得到的batch一般是長度為2的列表,列表的兩個值都是Tensor,分別表示數據(是一個batch的)和標簽。_get_batch()方法除了返回batch數據外,還得到另一個輸出:idx,這個輸出表示batch的index,這個if idx != self.rcvd_idx條件語句表示如果你讀取到的batch的index不等於當前想要的index:selg,rcvd_idx,那么就將讀取到的數據保存在字典self.reorder_dict中:self.reorder_dict[idx] = batch,然后繼續讀取數據,直到讀取到的數據的index等於self.rcvd_idx。

  DataloaderIter類的_get_batch方法。主要根據是否設置了超時時間來操作,如果超過指定的超時時間后沒有從隊列中讀到數據就報錯,如果不設置超時時間且一致沒有從隊列中讀到數據,那么就會一直卡着且不報錯,這部分是PyTorch后來修的一個bug。

  DataLoaderIter類的_process_next_batch方法。首先對self.rcvd_idx進行加一,也就是更新下下一個要讀取的batch數據的index。然后調用_put_indices()方法獲取下一個batch的每個數據的index。

  DataLoaderIter類的_put_indices方法。該方法主要實現從self.sample_iter中讀取下一個batch數據中每個數據的index:indices = next(self.sample_iter, None),注意這里的index和前面idx是不一樣的,這里的index是一個batch中每個數據的index,idx是一個batch的index;然后將讀取到的index通過調用queue對象的put方法壓到隊列self.index_queue中:self.index_queue.put((self.send_idx, indices))


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM