nn.MultiheadAttention

nn.MultiheadAttention是PyTorch中的一个模块，它实现了多头自注意力机制。在自然语言处理（NLP）中，自注意力机制是一种用于处理序列数据的技术，它能够捕捉到输入序列中不同位置之间的依赖关系，并且能够在不同的层次上进行抽象。

在nn.MultiheadAttention中，输入是一个三维的张量，形状为(batch_size, sequence_length, hidden_size)，其中hidden_size是每个时间步的向量维度。MultiheadAttention将输入张量拆分成多个头，在每个头上执行自注意力计算，最后将所有头的输出进行拼接和线性变换得到最终的输出。MultiheadAttention的输出形状与输入相同。

nn.MultiheadAttention的构造函数需要指定以下参数：

- embed_dim：输入张量的最后一个维度的大小
- num_heads：头的数量
- dropout：可选的dropout概率

除此之外，nn.MultiheadAttention还有一个重要的方法forward，它接受一个输入张量和一个可选的key张量和value张量，并返回该层的输出张量。在自然语言处理任务中，通常使用相同的输入张量作为key和value。