pytorch入门到放弃3：Dataset

pytorch为数据的读取和预处理提供了很多工具。

其中，torch.utils.data.Dataset是代表这一数据的抽象类，我们可以定义自己的数据class来继承和重写这个抽象类，只需要定义__len__和__getitem__两个函数就行，其中__getitem__包含了数据的读取和处理：

from torch.utils.data import Dataset, DataLoader
import pandas as pd


class MyDataset(Dataset):
    def __init__(self, csv_file, txt_file, root_dir):
        self.csv_data = pd.read_csv(csv_file)
        with open(txt_file, 'r') as f:
            self.txt_data = f.readlines()
        self.root_dir = root_dir


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

    def __getitem__(self, item):
        data = (self.csv_data[item], self.txt_data[item])
        return data

通过上面的方式，可以定义我们需要的类，然后通过torch.utils.data.DataLoader取得每一个数据，如下:

dataiter = DataLoader(MyDataset, batch_size=32,shuffle=True, num_workers=8)

创作不易，如果您觉的有用，可以来点赞助，一百两百不嫌多，一毛两毛不嫌少。