torch_collate_fn.py

def collate_fn(data):
    """
    default_collate比较好地实现了对图片等的操作，但是并不支持对文字等不等长序列的操作
    dataloader的一个参数
    输入data是list of (x, y), list的长度是batch_size

    返回xs, ys, lens 
    (batch_size, x或y的size)的tensor
    """
    # Sort a data list 
    data.sort(key=lambda x: len(x[1]), reverse=True)
    xs, ys = zip(*data)
    
    # Merge xs.
    xs = torch.stack(xs, 0)

    # Merge ys.
    lens = [len(e) for e in ys]
    ys = torch.zeros(len(ys), max(lens)).long()
    # padding
    for i, e in enumerate(ys):
        end = lens[i]
        ys[i, :end] = ys[:end] 

    return xs, ys, lens