pytorch自定义DataSet

例1

import torch.utils.data as Data


class TrainData(Data.Dataset):
    def __init__(self, data):
        self.data = data
        self.len = 10

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        # index的取值为[0, len]，len为__len__函数的返回值。
        # 一轮训练所取的训练数据的总数为len
        data = self.data
        # 返回为训练数据和标签，当然也可以根据需要选择返回值的个数
        return data[index], data[index]
a = torch.ones(100)
a = TrainData(a)
b = Data.DataLoader(a, batch_size=5)
for i, j in b:
    print(i.shape)

例2

class TrainData(Data.Dataset):
    def __init__(self):
        self.len = 10

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        # index的取值为[0, len]，len为__len__函数的返回值。
        # 一轮训练所取的训练数据的总数为len
        # 返回为训练数据和标签，当然也可以根据需要选择返回值的个数
        data = torch.randint(0, 10, (3,))
        return data[0], data[1], data[2]


a = TrainData()
b = Data.DataLoader(a, batch_size=5)
for i, j, k in b:
    print(i)

注意：如果在getitem方法中使用np.random，那么在不同epoch中得到的随机数序列很可能是相同的，因此，最好使用torch或random中的随机方法