PyTorch之Dataset和TensorDataset

Deep Learning系列 @cxx

Dataset v.s. TensorDataset

使用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数调成第一维。
这里真的要注意呀。