Is there a way to return batch image in the __get_item__ method of the dataloader?

chenchr · November 26, 2017, 11:58am

Hello. I found that the bottleneck of training procedure of my project to date is the data reading from disk. For a image with size 640480, I just need size of 320240 therefore I use the random_crop. However it will help if I can crop the sam image multi times and pack them up to a batch at the same time. As the get_item method should return tensor with shape CHW and them the dataloader pack up to NCHW, is there a way to pre pack up multi image in the get_item ? Thanks!

chenyuntc · November 26, 2017, 12:26pm

you can set batch_size = 1 in Dataloader, and write your dataset. You may refer to

github.com

pytorch/vision/blob/master/torchvision/datasets/folder.py#L123-128


def __getitem__(self, index):
    """
    Args:
        index (int): Index


    Returns:
        tuple: (image, target) where target is class_index of the target class.
    """
    path, target = self.imgs[index]
    img = self.loader(path)
    if self.transform is not None:
        img = self.transform(img)
    if self.target_transform is not None:
        target = self.target_transform(target)


    return img, target


def __len__(self):
    return len(self.imgs)

It’s something like:

   def __getitem__(self, index):

        path, target = self.imgs[index]
        img = self.loader(path)
        # ......
        for ii in range(batch_size):
             img_ = random_crop(img)
             imgs.append(img_)
        return torch.stack(imgs), target

chenchr · November 26, 2017, 12:36pm

thanks for you reply. However I need the batch size to be large than 1… And if I write torch.stack(imgs) in the getitem , the return data shape will be NCHW which should be CHW…

chenyuntc · November 26, 2017, 12:42pm

then you also need to write your collate_fn

github.com

pytorch/pytorch/blob/master/torch/utils/data/dataloader.py#L83


        raise
    if r is None:
        break
    if isinstance(r[1], ExceptionWrapper):
        out_queue.put(r)
        continue
    idx, batch = r
    try:
        if pin_memory:
            batch = pin_memory_batch(batch)
    except Exception:
        out_queue.put((idx, ExceptionWrapper(sys.exc_info())))
    else:
        out_queue.put((idx, batch))


numpy_type_map = {
'float64': torch.DoubleTensor,
'float32': torch.FloatTensor,
'float16': torch.HalfTensor,
'int64': torch.LongTensor,
'int32': torch.IntTensor,

something like

def my_collate(batch):
    imgs,targets = zip(*batch)
    return torch.cat(imgs),torch.cat(targets)

and use it in dataloader

dataloader = DataLoader(dataset,collate_fn=my_collate)

chenchr · November 26, 2017, 2:11pm

Thanks! I got your point.