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
,则Linear
和LayerNorm
层不会学习加性偏置。默认值:True
。
-
- 示例:
-
>>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8) >>> src = torch.rand(10, 32, 512) >>> out = encoder_layer(src)
-
或者,当
batch_first
为True
时: -
>>> 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_mode
或torch.no_grad
),要么传入的张量参数中没有设置requires_grad
-
训练已关闭(使用
.eval()
) -
当 batch_first 为
True
且输入数据是批量化的(即,src.dim() == 3
)时 -
激活函数可以是以下之一:\"relu\",\"gelu\",
torch.functional.relu
,或torch.functional.gelu
-
最多传递
src_mask
和src_key_padding_mask
中的一个 -
如果 src 是一个NestedTensor,则不会传递
src_mask
和src_key_padding_mask
。 -
两个
LayerNorm
实例的eps
值相同(除非调用者手动修改了一个实例而未修改另一个,否则这种情况会自然出现)
如果使用了优化实现,可以将NestedTensor 传递给
src
参数,以更高效地表示填充,而无需使用填充掩码。在这种情况下,将会返回一个 NestedTensor,并且可以期望获得与输入中填充部分比例成正比的额外加速。 -
- forward(src, src_mask=None, src_key_padding_mask=None, is_causal=False)[源代码]
-
将输入传递给编码器层。
- 参数
- 返回类型
- 形状:
-
参阅
Transformer
的文档。