PyTorch张量reshape函数用法解析：Transformer模型实现困惑

你在实现《Attention Is All You Need》论文的模型时，处理输入嵌入时遇到这段张量重塑代码，搞不清它的具体作用：

values = values.reshape(N, value_len, self.heads, self.head_dim)
#Here N, value_len and embed size is the dimension of the tensor "values"

这段代码是Transformer多头注意力机制的核心拆分步骤，具体作用如下：

• 原values张量的维度是(N, value_len, embed_size)，其中总嵌入维度embed_size必然等于self.heads * self.head_dim（头数 × 单头注意力的维度）
• 重塑操作把原本的单一大嵌入维度，拆分成多头+单头维度的结构，得到(N, value_len, 头数, 单头维度)的张量
• 拆分后，模型可以对每个头单独计算注意力，实现多头并行的注意力计算：每个头专注于输入序列中不同维度的语义信息，同时拆分低维度子嵌入也能降低单步计算的复杂度
• 后续完成多头注意力计算后，还会把这些头的结果拼接回原来的嵌入维度，完成完整的多头注意力流程

内容的提问来源于stack exchange，提问作者Anwesh saha