pytorch自定义TensorDataset

在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)