【Pytorch】简析DataLoader中的collate_fn参数

如在博文数据批量化处理类Dataset和DataLoader中所介绍的一样，DataLoader可通过collate_fn参数，对Dataset生成的mini-batch的可迭代数据进行进一步处理，而本文就简要介绍下该参数，并给出一个简单的例子。

1. collate_fn的设置、输入和输出

collate_fn应当是一个可调用对象，常见的可以是外部定义的函数或者lambda函数。其接受DataLoader不设置collate_fn参数时每个mini-batch输出的样本序列，所以可以理解为一个后处理的函数。

在不设置collate_fn参数时，DataLoader的mini-batch样本序列样式取决于对应dataset参数的设置。而dataset只要是覆写了__init__、__getitem__和__len__方法的Dataset子类即可，所以其输出形式可以很多样化。

而collate_fn对于上述的输出进行进一步的加工，可通过定义采样、切片、组合和数学操作等一系列操作，生成最终的训练数据。

其输出值一般建议为可直接为深度学习模型直接使用的张量对象，而大小规模为min-batch。

由上可知，整个处理过程可以理解为：

2. 简单示例

在博文示例的基础上，对其进行简单拓展，通过 collate_fn参数实现对两组输入的拼接。

import torch
from torch.utils.data import Dataset, DataLoader

A = torch.randn(128, 3)   # 待拼接的张量A
B = torch.randn(128, 2)   # 待拼接的张量B
C = torch.randn(128, 1)   #　labels张量

# 1. 利用Dataset封装数据集
class MyDataset(Dataset):
    def __init__(self, x1, x2, y):
        assert x1.size(0)==x2.size(0)==y.size(0)
        self.x1, self.x2, self.y = x1, x2, y
        
    def __getitem__(self, idx):
        return (self.x1[idx], self.x2[idx], self.y[idx])
    
    def __len__(self):
        return self.x1.size(0)

dataset = MyDataset(A, B, C)

# 2. 定义后处理的collate_fn函数
# 需要特别注意的是：输入的mini-batch输入，返回的为张量，因此要注意对格式进行转换和统一
def unify_fn(batch_data):
    x_ = [x1.tolist()+x2.tolist() for x1, x2, y in batch_data]
    y_ = [y.tolist() for x1, x2, y in batch_data]
    return torch.tensor(x_), torch.tensor(y_)
    
# 3. 利用DataLoader完成数据集的批量化
MyDataLoader = DataLoader(dataset=dataset, shuffle=True, batch_size=4, collate_fn=unify_fn)
for data_iter in MyDataLoader:
    x, y = data_iter
    # 进行训练