在pytorch中如果仅仅是训练数据和标签,完全可以使用TensorDataset进行构造
train_set = torch.utils.data.TensorDataset(torch.FloatTensor(self.train_X), torch.from_numpy(self.train_label).long())
train_loader = DataLoader(train_set, batch_size=batch_size, num_workers=0,shuffle=True)
但是这种做法有一个问题,就是无法对数据集进行自定义一些属性,所以我想自定义这个TensorDataset
import torch
from torch.utils.data import Dataset
class CustomTensorDataset(Dataset):
def __init__(self, dataset):
[data_X, data_y] = dataset
X_tensor, y_tensor = torch.tensor(data_X), torch.tensor(data_y)
#X_tensor, y_tensor = Tensor(data_X), Tensor(data_y)
tensors = (X_tensor, y_tensor)
assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors)
self.tensors = tensors
self.train_labels = y_tensor.long()
self.test_labels = y_tensor.long()
def __getitem__(self, index):
x = self.tensors[0][index]
y = self.tensors[1][index]
return x, y
def __len__(self):
return self.tensors[0].size(0)
train_dataset = CustomTensorDataset([X,y])
test_dataset = CustomTensorDataset([X,y])
print(train_dataset)
可以看到,这里的train_dataset是有属性的,例如tensors,train_labels,test_labels都是可以访问的
dataloader进行iter
import numpy as np
from torch.utils.data import TensorDataset,DataLoader
torch.manual_seed(1)
x = np.arange(100).reshape(20,5)
y = np.arange(20)
ds = TensorDataset(torch.from_numpy(x),torch.from_numpy(y))
dl= DataLoader(ds,batch_size=4,shuffle=True)
for idx,(x,y) in enumerate(dl):
print(idx)
print(x)
print(y)
print(x.shape)
print(y.shape)
break
结果如下
import numpy as np
from torch.utils.data import TensorDataset,DataLoader
torch.manual_seed(1)
x = np.arange(100).reshape(20,5)
y = np.arange(20)
ds = TensorDataset(torch.from_numpy(x),torch.from_numpy(y))
dl= DataLoader(ds,batch_size=4,shuffle=True)
iter_dl = iter(dl)
xx,yy = next(iter_dl)
print(xx)
print(yy)
print(xx.shape)
print(yy.shape)