确认PyTorch MultiHeadAttention中in_proj_weight的结构对应关系

结论是肯定的：in_proj_weight参数的前embed_dim个元素对应query，中间embed_dim个元素对应key，最后embed_dim个元素对应value。

从PyTorch的MultiHeadAttention源码逻辑可以验证这一点：

• 该层的输入投影操作通过_in_proj函数完成，in_proj_weight的总维度是3 * embed_dim × embed_dim（要同时完成q、k、v三个方向的投影）。
• 投影完成后，代码会将合并的投影结果按embed_dim维度切分为三个部分，切分顺序就是先取前embed_dim作为query的投影结果，中间embed_dim作为key的投影结果，最后embed_dim作为value的投影结果。
• 如果你手动拆分in_proj_weight，可以用q_weight, k_weight, v_weight = in_proj_weight.split(embed_dim, dim=0)，得到的三个张量就分别对应q、k、v的投影权重。

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