假设有一个保存为npy格式的numpy数据集,现在需要将其变为pytorch的数据集,并能够被数据加载器DataLoader所加载
首先自定义一个数据集类,继承torch.utils.data.Dataset类
在这个类中要实现__init__,__getitem__,__len__这三个方法,否则会报错
然后实例化这个类,得到train_data,最后将train_data放入DataLoader数据加载器,到此已经完成
注意,在下面这个代码中的x(也就是数据加载器加载出来的数据)的类型是tensor。也就是说,上面的实现中自动把numpy数据类型转化为了tensor类型
from torch.utils.data import Dataset, DataLoader
class MyDataset(Dataset):
"""
path:数据集存放路径
"""
def __init__(self, path):
self.data = np.load(path)
def __getitem__(self, index):
return self.data[index]
def __len__(self):
return len(self.data)
if __name__ == '__main__':
train_data = MyDataset(r"D:\dataset.npy")
load1 = DataLoader(train_data, batch_size=128, shuffle=True, pin_memory=True, num_workers=3)
for x in load1:
print(x.size())
有时候我们需要同时加载数据和其对应的标签,则需要将数据集和标签定义在同一个数据加载器中,这时可以采用以下方法:
from torch.utils.data import Dataset, DataLoader
class MyDataset(Dataset):
def __init__(self, data, label):
self.data = data
self.label = label
def __getitem__(self, index):
return self.data[index], self.label[index]
def __len__(self):
return len(self.label)
if __name__ == '__main__':
a = np.array([0,1,2,3,4,5])
b = np.array([6,7,8,9,10,11])
trainset = MyDataset(a, b)
train_loader = torch.utils.data.DataLoader(trainset, batch_size=2, shuffle=True, pin_memory=True,
num_workers=3)
for x, y in train_loader:
print(x, y)
重点在__getitem__方法的实现,需要同时返回数据和标签