Pytorch DataLoader doesn't return batched data

Viewed 909

My dataset is composed of image patches obtained from the original image (face patches and random outside of face patches). Patches are stored in a folder with a name of an original image from which patches originate. I created my own DataSet and DataLoader but when I iterate over the dataset data is not returned in batches. A batch of size 1 should include an array of tuples of patches and a label, so with the increased batch size, we should get an array of arrays of tuples with labels. But DataLoader returns only one array of tuples no matter the batch size.

My dataset:

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

Then I use it as such:

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

In this provided case the batch size is 10, so printing of the batch_Y returns the correct number (10). But the printing of the batch_X returns 9 which is number of patch pairs - returns only one sample from dataset instead of batch of 10 samples where each of them is length of 9.

1 Answers

You should return a one dimension higher tensor instead of a list of tensors in __get_item__ function call. You can use 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])
Related