import math import torch from torch import nn from d2l import torch as d2l num_hiddens, num_heads = 100, 5 attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,num_hiddens, num_heads, 0.5) attention.eval() MultiHeadAttention( (attention): DotProductAttention( (dropout): Dropout(p=0.5, inplace=False) ) (W_q): Linear(in_features=100, out_features=100, bias=False) (W_k): Linear(in_features=100, out_features=100, bias=False) (W_v): Linear(in_features=100, out_features=100, bias=False) (W_o): Linear(in_features=100, out_features=100, bias=False) ) batch_size, num_queries, num_kvpairs, valid_lens = 2, 4, 6, torch.tensor([3, 2]) X = torch.ones((batch_size, num_queries, num_hiddens)) # query（2，4，100） Y = torch.ones((batch_size, num_kvpairs, num_hiddens)) # key和value （2，6，100） output = attention(X, Y, Y, valid_lens) # 输出大小与输入的query的大小相同 output.shape torch.Size([2, 4, 100])

这段代码演示了使用之前定义的多头注意力机制 `MultiHeadAttention` 进行前向传播计算的例子。

首先，定义了隐藏层大小 `num_hiddens` 和头数 `num_heads`。然后，创建了一个 `MultiHeadAttention` 实例 `attention`，并传入相应的参数来初始化。

接下来，调用 `eval()` 方法将模型设为评估模式。然后，创建了输入张量 `X` 和 `Y`，分别表示查询序列和键值对序列，它们的形状分别为 `(batch_size, num_queries, num_hiddens)` 和 `(batch_size, num_kvpairs, num_hiddens)`。这里使用了全1的张量作为示例输入。

接着，调用 `attention` 实例进行前向传播计算，并传入输入张量 `X`、键和值都使用 `Y`，以及有效长度 `valid_lens`。最后，输出了计算结果的形状，即 `(batch_size, num_queries, num_hiddens)`。

可以看到，输出的形状与输入的查询序列的形状相同，说明多头注意力机制成功地对输入序列进行了处理，并得到了相应的输出。注意，在这个例子中，键和值序列的长度比查询序列长（`num_kvpairs > num_queries`），这是为了展示多头注意力机制能够处理不同长度的序列。