【Pytorch】Dataset_DataLoader

• 示例代码

import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import os


class TextDataset(Dataset):
    def __init__(self, data_path):
        data = np.loadtxt(data_path, delimiter=',')
        self._x = data[:, :2]
        self._y = data[:, -1]
        self.len = len(data)

    def __getitem__(self, idx):
        return self._x[idx], self._y[idx]

    def __len__(self):
        return self.len


class ImgDateset(Dataset):
    def __init__(self, root_dir):
        self.root_dir = root_dir
        self.img_path_list = os.listdir(self.root_dir)

    def __getitem__(self, idx):
        img_path = os.path.join(self.root_dir, self.img_path_list[idx])
        img = Image.open(img_path)
        # img.show()
        return img_path

    def __len__(self):
        return len(self.img_path_list)


if __name__ == '__main__':
    i = 1
    if i:
        text_dataset = TextDataset('./example_data.txt')
        TextDataLoader = DataLoader(text_dataset, batch_size=3, shuffle=True, drop_last=False)

        iteration = iter(TextDataLoader)
        while True:
            try:
                xs, ys = next(iteration)
                print('Value :', xs)
                print('Table:', ys)
            except StopIteration:
                break
    else:
        img_dataset = ImgDateset('./img_data')
        img = img_dataset.__getitem__(2)
        print(img_dataset)

        ImgDataLoader = DataLoader(dataset=img_dataset, batch_size=3, shuffle=False, drop_last=False)
        print(len(ImgDataLoader))
        for img_path in ImgDataLoader:
            print(img_path)

• example_data.txt

1.0,12.0,12.0
2.0,23.4,23.5
3.0,34.5,34.6
4.0,45.6,45.7
5.0,56.7,56.8
6.0,67.8,67.9
7.0,78.9,78.0
8.0,89.0,89.1

• img_data