如何在 transformer 编码器中发送注意掩码“掩码”矩阵以及我在 pytorch 的 nn.TransformerEncoder 中的潜在值？

Question

我有一个 transformer_encoder 在 pytorch 中初始化为：

transformer_encoder = nn.TransformerEncoder(nn.TransformerEncoderLayer(d_model=input_dim, nhead=num_heads), num_layers=num_layers)

现在我有一个潜在的，它的形状是 [8, 320, 512]，其中 8 是 batch_size，320 是序列长度，512 是嵌入维度。

我还有一个形状为 [320, 320] 的注意力面具。

我想将它传递给变压器编码器，例如：

latent = torch.rand(8, 320, 512)
mask = torch.rand(320, 320)

output = transformer_encoder(latent, mask)

但它给我错误：
“2D attn_mask 的形状是 torch.Size([320, 320])，但应该是 (8, 8)。”

现在，我沿着批次维度重复矩阵，创建其中的 8 个 (320, 320) 矩阵，形成“掩码”的形状，[8, 320, 320]。当我再次运行上面的代码

latent.shape = [8, 320, 512]
mask.shape = [8, 320, 320]

output = transformer_encoder(latent, mask)

它给了我这个错误：
“3D attn_mask 的形状是 torch.Size([8, 320, 320])，但应该是 (2560, 8, 8)。”

有人可以解释这里的问题是什么吗？我该如何解决？或者正确的实现方式？

Answer 1

您的面罩未正确设置。你的批量大小为 8，序列长度为 320。尝试添加：

mask = torch.ones(8, 320, 320)
mask = mask.masked_fill(mask == 0, float('-inf'))

使用前：

output = transformer_encoder(latent, mask)