Pytorch collate_fn用法

By default, Dataloader use collate_fn method to pack a series of images and target as tensors (first dimension of tensor is batch size). The default collate_fn expects all the images in a batch to have the same size because it uses torch.stack() to pack the images. If the images provided by Dataset have variable size, you have to provide your custom collate_fn. A simple example is shown below:

 # a simple custom collate function, just to show the idea

 # `batch` is a list of tuple where first element is image tensor and

 # second element is corresponding label

 def my_collate(batch):
     data = [item[0] for item in batch]  # just form a list of tensor

     target = [item[1] for item in batch]
     target = torch.LongTensor(target)
     return [data, target]

Reference:   Writing Your Own Custom Dataset for Classification in PyTorch

By default, torch stacks the input image to from a tensor of size N*C*H*W, so every image in the batch must have the same height and width. In order to load a batch with variable size input image, we have to use our own collate_fn which is used to pack a batch of images.

For image classification, the input to collate_fn is a list of with size batch_size. Each element is a tuple where the first element is the input image(a torch.FloatTensor) and the second element is the image label which is simply an int. Because the samples in a batch have different size, we can store these samples in a list ans store the corresponding labels in torch.LongTensor. Then we put the image list and the label tensor into a list and return the result.

here is a very simple snippet to demonstrate how to write a custom collate_fn:

 import torch
 from torch.utils.data import DataLoader
 from torchvision import transforms
 import torchvision.datasets as datasets
 import matplotlib.pyplot as plt

 # a simple custom collate function, just to show the idea
 def my_collate(batch):
     data = [item[0] for item in batch]
     target = [item[1] for item in batch]
     target = torch.LongTensor(target)
     return [data, target]

 def show_image_batch(img_list, title=None):
     num = len(img_list)
     fig = plt.figure()
     for i in range(num):
         ax = fig.add_subplot(1, num, i+1)
         ax.imshow(img_list[i].numpy().transpose([1,2,0]))
         ax.set_title(title[i])

     plt.show()

 #  do not do randomCrop to show that the custom collate_fn can handle images of different size
 train_transforms = transforms.Compose([transforms.Scale(size = 224),
                                        transforms.ToTensor(),
                                        ])

 # change root to valid dir in your system, see ImageFolder documentation for more info
 train_dataset = datasets.ImageFolder(root="/hd1/jdhao/toyset",
                                      transform=train_transforms)

 trainset = DataLoader(dataset=train_dataset,
                       batch_size=4,
                       shuffle=True,
                       collate_fn=my_collate, # use custom collate function here
                       pin_memory=True)

 trainiter = iter(trainset)
 imgs, labels = trainiter.next()

 # print(type(imgs), type(labels))
 show_image_batch(imgs, title=[train_dataset.classes[x] for x in labels])

Reference:    How to create a dataloader with variable-size input

Dataloader的测试用例：

 import torch
 import torch.utils.data as Data
 import numpy as np

 test = np.array([0,1,2,3,4,5,6,7,8,9,10,11])

 inputing = torch.tensor(np.array([test[i:i + 3] for i in range(10)]))
 target = torch.tensor(np.array([test[i:i + 1] for i in range(10)]))

 torch_dataset = Data.TensorDataset(inputing,target)
 batch = 3

 loader = Data.DataLoader(
     dataset=torch_dataset,
     batch_size=batch, # 批大小
     # 若dataset中的样本数不能被batch_size整除的话，最后剩余多少就使用多少
     collate_fn=lambda x:(
         torch.cat(
             [x[i][j].unsqueeze(0) for i in range(len(x))], 0
             ).unsqueeze(0) for j in range(len(x[0]))
         )
     )

 for (i,j) in loader:
     print(i)
     print(j)

Reference： DataLoader的collate_fn参数

pytorch 读取变长数据

https://zhuanlan.zhihu.com/p/60129684