Transformer编码层

classtorch.nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation=<function relu>, layer_norm_eps=1e-05, batch_first=False, norm_first=False, bias=True, device=None, dtype=None)[源代码]

TransformerEncoderLayer 包含自注意力机制和前向网络。

此标准编码器层基于论文“Attention Is All You Need”,作者是 Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser 和 Illia Polosukhin。该论文发表于 2017 年的 Advances in Neural Information Processing Systems,页码为 6000-6010。用户在应用过程中可能会对此进行修改或以不同方式实现。

TransformerEncoderLayer 可以处理传统的 torch.tensor 输入或 Nested Tensor 输入。派生类也应该同时支持这两种输入格式。(由于 Nested Tensor 仍处于原型阶段,目前并非所有的输入组合都能被 TransformerEncoderLayer 支持。)

如果你在实现一个自定义层,可以从中选择从Module类或TransformerEncoderLayer类进行派生。如果自定义层同时支持torch.Tensors和Nested Tensors输入,请将其实现为TransformerEncoderLayer的子类。如果自定义层仅支持torch.Tensor输入,则使其继承自Module。

参数
  • d_model (int) – 输入数据中预期的特征数目(必填)。

  • nhead (int) – 多头注意力模型中头部的数量(必填)。

  • dim_feedforward (int) – 前向网络模型的维度,默认值为2048。

  • dropout (浮点数) - Dropout值,默认为0.1。

  • activation (Union[str, Callable[[Tensor], Tensor]]) – 中间层的激活函数,可以是字符串(如“relu”或“gelu”),也可以是一元可调用函数。默认值:relu

  • layer_norm_eps (float) – 层归一化组件中的 eps 值,默认为 1e-5。

  • batch_first (bool) – 如果为 True,则输入和输出张量的形状为 (批次, 序列, 特征)。默认值:False(序列, 批次, 特征)。

  • norm_first (bool) – 如果为 True,则在注意力和前馈操作之前进行层归一化。否则在之后进行。默认值:False(在之后)。

  • bias (bool) – 如果设置为False,则LinearLayerNorm 层不会学习加性偏置。默认值:True

示例:
>>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
>>> src = torch.rand(10, 32, 512)
>>> out = encoder_layer(src)
或者,当 batch_firstTrue 时:
>>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
>>> src = torch.rand(32, 10, 512)
>>> out = encoder_layer(src)
快捷路径:

如果满足以下所有条件,forward() 将使用在 FlashAttention: 快速且内存高效的精确注意机制与 IO 感知 中描述的特殊优化实现:

  • 自动求梯度要么被禁用了(使用了torch.inference_modetorch.no_grad),要么传入的张量参数中没有设置requires_grad

  • 训练已关闭(使用.eval()

  • 当 batch_first 为 True 且输入数据是批量化的(即,src.dim() == 3)时

  • 激活函数可以是以下之一:\"relu\",\"gelu\",torch.functional.relu,或 torch.functional.gelu

  • 最多传递 src_masksrc_key_padding_mask 中的一个

  • 如果 src 是一个NestedTensor,则不会传递 src_masksrc_key_padding_mask

  • 两个LayerNorm 实例的eps 值相同(除非调用者手动修改了一个实例而未修改另一个,否则这种情况会自然出现)

如果使用了优化实现,可以将NestedTensor 传递给 src 参数,以更高效地表示填充,而无需使用填充掩码。在这种情况下,将会返回一个 NestedTensor,并且可以期望获得与输入中填充部分比例成正比的额外加速。

forward(src, src_mask=None, src_key_padding_mask=None, is_causal=False)[源代码]

将输入传递给编码器层。

参数
  • src (Tensor) – 输入到编码层的序列(必填)。

  • src_mask (Optional[Tensor]) – 源序列的可选掩码。

  • src_key_padding_mask (Optional[Tensor]) – 源键的批掩码(可选)。

  • is_causal (bool) – 如果指定,则将因果掩码应用于 src mask。默认值为False。警告:如果 is_causal 设置为真,表示 src_mask 是因果掩码。提供错误的提示可能导致执行出错,包括前向和后向兼容性问题。

返回类型

Tensor

形状:

参阅Transformer的文档。

本页目录