pytorch dataloader修改shape的问题

不管是list类型的数据还是dict类型的数据，输入给dataloader都会把我的数据shape给转置了。

比如我的数据应该是[batch_size, seq_len]尺寸的，dataloader吐出来的却是[seq_len, batch_size]。如果还是tensor类型还行，但是它会变成List[Tensor]类型。如果只是两维也还行，但是如果是[batch_size, seq_len, nvars]呢，很难把shape变回来。

之前其实没有遇到过这个问题。通过查看源码，加上和之前的代码对比，终于找到了原因：

    elif isinstance(elem, collections.abc.Sequence):
        # check to make sure that the elements in batch have consistent size
        it = iter(batch)
        elem_size = len(next(it))
        if not all(len(elem) == elem_size for elem in it):
            raise RuntimeError('each element in list of batch should be of equal size')
        transposed = list(zip(*batch))  # It may be accessed twice, so we use a list.

dataloader默认的collate_fn是这样处理的。也就是说只要是collections.abs.Sequence类型的，都会被转置。

那么解决办法就明显了，只要数据不是这个类型的就行。

所以在准备dataset的时候要把数据转成numpy或者tensor格式，这样就不会走转置这个路径了。

如果用的是huggingface的datasets，可以用dataset.set_format(type='numpy')来设置数据格式，且它只对数字类型的起作用，不会影响文本类型，可以放心使用。

实在搞不懂collate默认转置的意义是什么。开发者的解释在这里default_collate: if elem in batch is list, the whole batch would be transpose? · Issue #50272 · pytorch/pytorch (github.com)

感兴趣的可以看下。