PyTorch中BucketIterator输出张量维度为何与预期相反？

你看到的反序形状完全正常，这是PyTorch针对序列模型的标准设计，不是错误。

核心原因：循环模型的输入格式要求 #

PyTorch里的RNN、LSTM、GRU这类循环神经网络层，默认输入格式是(序列长度, batch_size, 特征维度)。BucketIterator作为torchtext里专门为序列任务设计的迭代器，直接输出了符合这个要求的张量形状——也就是你看到的(seq_len, 64)，而非你预期的(64, seq_len)。

这么设计是因为：循环模型会按时间步（序列的每个token位置）依次处理整个batch的样本。比如序列长度是15，模型先处理64个样本的第1个token，再处理所有样本的第2个token，以此类推。这种按时间步批量处理的方式，能让循环层的计算更高效，也符合循环模型的运行逻辑。

如何得到你预期的(64, 序列长度)形状？ #

如果你更习惯用批量优先的格式，只要在创建BucketIterator时加上batch_first=True参数就行：

train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
    (train_data, valid_data, test_data),
    batch_size=batch_size,
    sort_within_batch=True,
    sort_key=lambda x: len(x.src),
    device=device,
    batch_first=True  # 添加该参数
)

此时输出的张量形状就会变成(64, seq_len)。但要注意：如果你的模型用到了循环层，必须同时在循环层初始化时设置batch_first=True，否则模型会报错。比如：

lstm = nn.LSTM(input_size=emb_dim, hidden_size=hid_dim, batch_first=True)

额外提醒：Transformer模型的差异 #

要是后续你用Transformer模型，它的默认输入格式就是批量优先的(64, seq_len, 特征维度)，这时候用batch_first=True的迭代器输出就刚好匹配，不用额外调整。

内容的提问来源于stack exchange，提问作者Zaharie Andrei