你所不知道的 Pytorch 大補包(二)：Dataset DataLoader

發表於 更新於 分類於 Disqus： 文章字數： 1.1k 所需閱讀時間 ≈ 3 分鐘

如果今天開發需求不是像 Mnist 這樣，別人已經幫你準備好的資料集，而是自己的影像資料集，那要怎麼放進 DataLoader 裡面訓練呢？

keywords: DataLoader

Dataset DataLoader

  • 使用繼承 Dataset 可以自定義 data，再放進 DataLoader 中
  • 一個 Dataset 繼承後要 override 的 function 如下：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
from torch.utils.data.dataset import Dataset

class customDataset(Dataset):
    def __init__(self):
        # --------------------------------------------
        # Initialize paths, transforms, and so on
        # --------------------------------------------
        
    def __getitem__(self, index):
        # --------------------------------------------
        # 1. Read from file (using numpy.fromfile, PIL.Image.open)
        # 2. Preprocess the data (torchvision.Transform).
        # 3. Return the data (e.g. image and label)
        # --------------------------------------------
        
    def __len__(self):
        # --------------------------------------------
        # Indicate the total size of the dataset
        # --------------------------------------------
  • __init__ 負責初使化 path、img list、label list、transform
  • __getitem__ 負責讀取圖片，並做 transform，回傳 img 以及 label
    • 回傳值也可以回傳不只兩個，也可根據需求回傳想要的資料
  • __len__ 回傳 imgs 的長度 len(self.imgs)