PyTorch之Dataset和TensorDataset

Deep Learning系列 @cxx

使用PyTorch搭建过Neural Network的小伙伴们都知道，在数据准备步骤里，我们需要把训练集的x和y分装在dataset里，然后将dataset分装到DataLoader中去，便于之后在搭建好的模型中训练。
简言之，dataset是用来做打包和预处理（比如输入资料路径自动读取）；DataLoader则是将整个资料集（dataset）按照batch进行迭代分装或者shuffle（可以得到一个iterator以利于for循环读取）。

Dataset

如果使用继承Dataset的方式，那么在自定义的dataset类中必须给予__len__和__getitem__的定义。
进行图片处理的时候，可以定义一个transforms来随机旋转训练图片，将图片格式变成tensor等
（这里有一个坑）

假设我们读取了一个有如下格式的图片

将图片分装到dataset里，再放到dataloader里

from torch.utils.data import TensorDataset
batch_size = 128
train_transform = transforms.Compose([
transforms.ToPILImage(),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ToTensor(),]
)
test_transform = transforms.Compose(
[transforms.ToPILImage(),
transforms.ToTensor(),]
) #测试集不需要翻转或旋转图片
#继承Dataset
class ImgDataset(Dataset):
def __init__(self, x, y=None, transform=None):
self.x = x
self.y = y
# label is required to be a LongTensor
if y is not None:
self.y = torch.LongTensor(y)
self.transform = transform
def __len__(self):
return len(self.x)
def __getitem__(self, index):
X = self.x[index]
if self.transform is not None:
X = self.transform(X)
if self.y is not None:
Y = self.y[index]
return X, Y
else:
return X


#将dataset分装到dataloader里
train_dataloader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True
)
test_dataloader = DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False
)


接下来我们可以输出一个batch看看图片的格式

我们发现一个batch的x[0]的shape由原先的(128, 128, 3)变成了(3, 128, 128)。
原因在于transformers.toTensor()方法有自动转换维度的功能，它会将channel变成第一维（夺么坑爹的功能，导致我排查了好久不知道是哪里出了问题==）
具体可以参照这篇博客transforms.ToTensor()本身有维度转换功能

TensorDataset

张量资料集tensrdataset是最常见的形式，因为PyTorch本身有提供方便的TensorDataset给我们使用

torch.utils.data.TensorDataset(data_tensor, target_tensor)


用TensorDataset写会少写很多东西

#将资料转换成tensor
tsr_x_train, tsr_y_train = torch.tensor(x_train), torch.tensor(y_train)
tsr_x_val, tsr_y_val = torch.tensor(x_val), torch.tensor(y_val)
tsr_x_testing = torch.tensor(x_test)
#然后只需要一行就可以啦
train_dataset = TensorDataset(tsr_x_train, tsr_y_train)
val_dataset = TensorDataset(tsr_x_val, tsr_y_val)
#装入dataloader的步骤同上
train_dataloader = DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True
)
test_dataloader = DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False
)


我们跑一个loop看看这次维度是否被转换了

答案是：这次没有！
这次的x[0]的shape同我们一开始设置的shape，TensorDataset并没有帮我们把channel数调成第一维。
这里真的要注意呀。