Pytorch DataLoader 没有 return 批处理数据

Question

我的数据集由从原始图像获得的图像块组成（面部块和面部块外的随机块）。补丁存储在一个文件夹中，该文件夹的名称与补丁所源自的原始图像相同。我创建了自己的 DataSet 和 DataLoader，但是当我遍历数据集时，数据并没有批量返回。大小为 1 的批次应该包括一个补丁元组数组和一个标签，因此随着批次大小的增加，我们应该得到一个包含标签的元组数组数组。但是 DataLoader returns 无论批大小如何，都只有一个元组数组。

我的数据集：

  import os
  import cv2 as cv
  import PIL.Image as Image
  import torchvision.transforms as Transforms
  from torch.utils.data import dataset    

  class PatchDataset(dataset.Dataset):
    def __init__(self, img_folder, n_patches):
      self.img_folder = img_folder
      self.n_patches = n_patches
      self.img_names = sorted(os.listdir(img_folder))

      self.transform = Transforms.Compose([
        Transforms.Resize((50, 50)),
        Transforms.ToTensor()
      ])
    
    def __len__(self):
      return len(self.img_names)
    
    def __getitem__(self, idx):
      img_name = self.img_names[idx]
      patch_dir = os.path.join(self.img_folder, img_name)
      patches = []
    
      for i in range(self.n_patches):
        face_patch = cv.imread(os.path.join(patch_dir, f'{str(i)}_face.png'))
        face_patch = cv.cvtColor(face_patch, cv.COLOR_BGR2RGB)
        face_patch = Image.fromarray(face_patch)
        face_patch = self.transform(face_patch)
    
        patch = cv.imread(os.path.join(patch_dir, f'{str(i)}_patch.png'))
        patch = cv.cvtColor(patch, cv.COLOR_BGR2RGB)
        patch = Image.fromarray(patch)
        patch = self.transform(patch)

        patches.append((face_patch, patch))
    
      return patches, int(img_name.split('-')[0])

然后我就这样使用它：

X = PatchDataset(PATCHES_DIR, 9)
train_dl = dataloader.DataLoader(
    X,
    batch_size=10,
    drop_last=True
)

for batch_X, batch_Y in train_dl:
  print(len(batch_X))
  print(len(batch_Y))

在本例中，批量大小为 10，因此 batch_Y returns 的打印是正确的数字 (10)。但是 batch_X returns 9 的打印是补丁对的数量 - returns 只有数据集中的一个样本，而不是 10 个样本的批次，其中每个样本的长度为 9.

Answer 1

在 __get_item__ 函数调用中，您应该 return 高一维 tensor 而不是 list 张量。您可以使用 torch.stack(patches).

def __getitem__(self, idx):
   img_name = self.img_names[idx]
   patch_dir = os.path.join(self.img_folder, img_name)
   patches = []

   for i in range(self.n_patches):
       face_patch = cv.imread(os.path.join(patch_dir, f'{str(i)}_face.png'))
       face_patch = cv.cvtColor(face_patch, cv.COLOR_BGR2RGB)
       face_patch = Image.fromarray(face_patch)
       face_patch = self.transform(face_patch)

       patch = cv.imread(os.path.join(patch_dir, f'{str(i)}_patch.png'))
       patch = cv.cvtColor(patch, cv.COLOR_BGR2RGB)
       patch = Image.fromarray(patch)
       patch = self.transform(patch)

       patches.append((face_patch, patch))

   return torch.stack(patches), int(img_name.split('-')[0])

Pytorch DataLoader 没有 return 批处理数据

Pytorch DataLoader doesn't return batched data

python

dataset

pytorch

pytorch-dataloader