Python PyTorch Transformer.forward用法及代码示例

本文简要介绍python语言中 torch.nn.Transformer.forward 的用法。

用法:

forward(src, tgt, src_mask=None, tgt_mask=None, memory_mask=None, src_key_padding_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None)

参数：

• src-编码器的序列(必需)。

• tgt-解码器的序列(必需)。

• src_mask-src 序列的附加掩码(可选)。

• tgt_mask-tgt 序列的附加掩码(可选)。

• memory_mask-编码器输出的附加掩码(可选)。

• src_key_padding_mask-每批 src key 的 ByteTensor 掩码(可选)。

• tgt_key_padding_mask-每批 tgt 键的 ByteTensor 掩码(可选)。

• memory_key_padding_mask-每个批次的内存 key 的ByteTensor 掩码(可选)。

接收并处理屏蔽的源/目标序列。

形状：

• src: (S, N, E) , (N, S, E) 如果 batch_first。

• tgt: (T, N, E) , (N, T, E) 如果 batch_first。

• src_mask：(S, S)。

• tgt_mask：(T, T)。

• memory_mask：(T, S)。

• src_key_padding_mask：(N, S)。

• tgt_key_padding_mask：(N, T)。

• memory_key_padding_mask：(N, S)。

注意：[src/tgt/memory]_mask 确保允许位置 i 参加未屏蔽的位置。如果提供ByteTensor，则非零位不得参加，零位不变。如果提供了BoolTensor，则True 的位置不允许参加，而False 的值将保持不变。如果提供了FloatTensor，它将被添加到注意力权重中。 [src/tgt/memory]_key_padding_mask提供了key中的指定元素被注意忽略。如果提供了ByteTensor，则将忽略非零位置，而零位置将保持不变。如果提供了BoolTensor，则将忽略具有True 值的位置，而具有False 值的位置将保持不变。

• 输出：(T, N, E)，如果batch_first，则为(N, T, E)。

注意：由于 Transformer 模型中的multi-head attention 架构，transformer 的输出序列长度与解码的输入序列(即目标)长度相同。

其中 S 是源序列长度，T 是目标序列长度，N 是批量大小，E 是特征数

例子

>>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)