解析torch官方代码脚本文件:transformer.py。版本:1.9.1+cu111。
首先查看《Torch中多头注意力MultiheadAttention
的中文注释》解析;
最后查看下方transformer解析。
话不多说,看代码吧!
import copy
from typing import Optional, Any
import torch
from torch import Tensor
from .. import functional as F
from .module import Module
from .activation import MultiheadAttention
from .container import ModuleList
from ..init import xavier_uniform_
from .dropout import Dropout
from .linear import Linear
from .normalization import LayerNorm
class Transformer(Module):
r"""这是一个变换器模型,用户可以根据需要修改其属性。该架构基于论文《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页。用户可以使用对应的参数构建BERT模型(参见https://arxiv.org/abs/1810.04805)。
参数如下:
d_model:编码器/解码器输入中预期的特征数量(默认为512)。
nhead:多头注意力模型中的头数(默认为8)。
num_encoder_layers:编码器中子编码层的数量(默认为6)。
num_decoder_layers:解码器中子解码层的数量(默认为6)。
dim_feedforward:前馈网络模型的维度(默认为2048)。
dropout:丢弃率(默认为0.1)。
activation:编码器/解码器中间层的激活函数,可选relu或gelu(默认为relu)。
custom_encoder:自定义编码器(默认为None)。
custom_decoder:自定义解码器(默认为None)。
layer_norm_eps:层归一化组件中的eps值(默认为1e-5)。
batch_first:如果设为True,则输入和输出张量将以(批次,序列,特征)的形式提供。默认值为False(序列,批次,特征)。
Examples::
>>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
>>> src = torch.rand((10, 32, 512))
>>> tgt = torch.rand((20, 32, 512))
>>> out = transformer_model(src, tgt)
Note: A full example to apply nn.Transformer module for the word language model is available in
https://github.com/pytorch/examples/tree/master/word_language_model
__init__:
batch_first: 如果为 True,则输入和输出张量将按 (batch, seq, feature) 的顺序提供;否则,顺序为 (seq, batch, feature)。
编码器和解码器组件:
self.encoder: 如果提供了 custom_encoder,则直接使用这个自定义编码器;
否则,创建一个标准的 TransformerEncoder 实例,它由多个 TransformerEncoderLayer 层堆叠构成。
self.decoder: 同样地,如果提供了 custom_decoder,则直接使用;
否则,创建一个标准的 TransformerDecoder 实例,由多个 TransformerDecoderLayer 层构成。
"""
def __init__(self, d_model: int = 512, nhead: int = 8, num_encoder_layers: int = 6,
num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,
activation: str = "relu", custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None,
layer_norm_eps: float = 1e-5, batch_first: bool = False,
device=None, dtype=None) -> None:
factory_kwargs = {'device': device, 'dtype': dtype}
super(Transformer, self).__init__()
if custom_encoder is not None:
self.encoder = custom_encoder
else:
encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout,
activation, layer_norm_eps, batch_first,
**factory_kwargs)
encoder_norm = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
if custom_decoder is not None:
self.decoder = custom_decoder
else:
decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout,
activation, layer_norm_eps, batch_first,
**factory_kwargs)
decoder_norm = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
self._reset_parameters()
self.d_model = d_model
self.nhead = nhead
self.batch_first = batch_first
def forward(self, src: Tensor, tgt: Tensor, src_mask: Optional[Tensor] = None, tgt_mask: Optional[Tensor] = None,
memory_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None,
tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
r"""接受并处理带有掩码的源序列和目标序列。
参数:
src: 输入到编码器的序列(必需)。
tgt: 输入到解码器的序列(必需)。
src_mask: 源序列的加性掩码(可选)。
tgt_mask: 目标序列的加性掩码(可选)。
memory_mask: 编码器输出的加性掩码(可选)。
src_key_padding_mask: 每批源键的ByteTensor掩码(可选)。
tgt_key_padding_mask: 每批目标键的ByteTensor掩码(可选)。
memory_key_padding_mask: 每批记忆键的ByteTensor掩码(可选)。
形状:
- src: :math:`(S, N, E)`,如果`batch_first`为真,则为`(N, S, E)`。
- tgt: :math:`(T, N, E)`,如果`batch_first`为真,则为`(N, T, E)`。
- src_mask: :math:`(S, S)`。
- tgt_mask: :math:`(T, T)`。
- memory_mask: :math:`(T, S)`。
- src_key_padding_mask: :math:`(N, S)`。
- tgt_key_padding_mask: :math:`(N, T)`。
- memory_key_padding_mask: :math:`(N, S)`。
注意: [src/tgt/memory]_mask确保位置i能访问未被掩码的位置。如果提供的是ByteTensor,非零位置不允许访问,而零位置保持不变。如果提供的是BoolTensor,`True`的位置不允许访问,而`False`值将保持不变。如果提供的是FloatTensor,它将被添加到注意力权重中。
[src/tgt/memory]_key_padding_mask提供了在键中应被注意力忽略的指定元素。如果提供的是ByteTensor,非零位置将被忽略,而零位置保持不变。如果提供的是BoolTensor,值为`True`的位置将被忽略,而值为`False`的位置将保持不变。
- 输出: :math:`(T, N, E)`,如果`batch_first`为真,则为`(N, T, E)`。
注意: 由于变换器模型中的多头注意力架构,变换器的输出序列长度与输入序列(即解码的目标)的长度相同。
其中S是源序列长度,T是目标序列长度,N是批次大小,E是特征数量。
Examples:
>>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)
"""
if not self.batch_first and src.size(1) != tgt.size(1):
raise RuntimeError("the batch number of src and tgt must be equal")
elif self.batch_first and src.size(0) != tgt.size(0):
raise RuntimeError("the batch number of src and tgt must be equal")
if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
raise RuntimeError("the feature number of src and tgt must be equal to d_model")
memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
tgt_key_padding_mask=tgt_key_padding_mask,
memory_key_padding_mask=memory_key_padding_mask)
return output
def generate_square_subsequent_mask(self, sz: int) -> Tensor:
r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
Unmasked positions are filled with float(0.0).
"""
mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
return mask
def _reset_parameters(self):
r"""Initiate parameters in the transformer model."""
for p in self.parameters():
if p.dim() > 1:
xavier_uniform_(p)
class TransformerEncoder(Module):
r"""TransformerEncoder is a stack of N encoder layers
参数:
- encoder_layer: TransformerEncoderLayer类的一个实例(必需)。
- num_layers: 编码器中子编码器层的数量(必需)。
- norm: 层归一化组件(可选)。
__init__:
encoder_layer: 这是单个编码器层的实例,通常由 TransformerEncoderLayer 构造。
num_layers: 表示要堆叠的编码器层数量。
norm: 可选参数,用于指定层归一化组件。
在初始化过程中,通过 _get_clones 函数复制 encoder_layer 指定次数来创建编码器层的列表,这个列表存储在 self.layers 中。self.num_layers 存储了编码器层数量,而 self.norm 则保存了提供的层归一化组件(如果有)
Examples::
>>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
>>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
>>> src = torch.rand(10, 32, 512)
>>> out = transformer_encoder(src)
"""
__constants__ = ['norm']
def __init__(self, encoder_layer, num_layers, norm=None):
super(TransformerEncoder, self).__init__()
self.layers = _get_clones(encoder_layer, num_layers)
self.num_layers = num_layers
self.norm = norm
def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
r"""依次通过编码器层传递输入.
参数:
Src:给编码器的序列(必需的)。
Mask: SRC序列的掩码(可选)。
Src_key_padding_mask:每批SRC键的掩码(可选)。
形状:
请参阅Transformer类中的文档。
流程:
src: 输入序列,是必须提供的参数。
mask: 可选参数,用于输入序列的掩码。
src_key_padding_mask: 可选参数,用于每批次源键的掩码。
在前向传播中,数据 src 依次通过每一个编码器层,其中可以使用 mask 和 src_key_padding_mask 来控制哪些部分需要被忽略。最后,如果提供了层归一化组件 norm,那么整个编码器的输出将会通过该组件进行归一化处理
"""
output = src
for mod in self.layers:
output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
if self.norm is not None:
output = self.norm(output)
return output
class TransformerDecoder(Module):
r"""TransformerDecoder 是由 N 个解码器层堆叠而成的模块
参数:
decoder_layer: 必须提供的 TransformerDecoderLayer() 类的一个实例。这定义了单个解码器层的行为和结构。
num_layers: 解码器中包含的子解码器层数量,这是一个必需的参数,决定了解码器的深度。
norm: 层归一化组件,这是一个可选参数,用于在解码器所有层的输出之后执行归一化操作,有助于稳定训练过程并加速收敛。
Examples::
>>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
>>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
>>> memory = torch.rand(10, 32, 512)
>>> tgt = torch.rand(20, 32, 512)
>>> out = transformer_decoder(tgt, memory)
"""
__constants__ = ['norm']
def __init__(self, decoder_layer, num_layers, norm=None):
super(TransformerDecoder, self).__init__()
self.layers = _get_clones(decoder_layer, num_layers)
self.num_layers = num_layers
self.norm = norm
def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
r"""依次通过解码器层传递输入(和掩码)。
参数:
Tgt:到解码器的序列(必需)。
存储器:来自编码器最后一层的序列(必需)。
Tgt_mask: TGT序列的掩码(可选)。
Memory_mask:内存序列的掩码(可选)。
Tgt_key_padding_mask:每批TGT键的掩码(可选)。
Memory_key_padding_mask:每批内存键的掩码(可选)。
形状:
请参阅Transformer类中的文档。
"""
output = tgt
for mod in self.layers:
output = mod(output, memory, tgt_mask=tgt_mask,
memory_mask=memory_mask,
tgt_key_padding_mask=tgt_key_padding_mask,
memory_key_padding_mask=memory_key_padding_mask)
if self.norm is not None:
output = self.norm(output)
return output
class TransformerEncoderLayer(Module):
r"""`TransformerEncoderLayer` 由自注意力(self-attn)和前馈网络(feedforward network)组成。
这个标准的编码器层基于论文 "Attention Is All You Need"。
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
Lukasz Kaiser, 和 Illia Polosukhin。2017年。注意力就是你所需要的。在《神经信息处理系统进展》期刊中,
第6000至6010页。用户在应用过程中可以对其进行修改或以不同方式实现。
参数:
d_model: 输入中预期特征的数量(必需)。
nhead: 多头注意力模型中的头数(必需)。
dim_feedforward: 前馈网络模型的维度(默认=2048)。
dropout: dropout的值(默认=0.1)。
activation: 中间层的激活函数,可选relu或gelu(默认=relu)。
layer_norm_eps: 层归一化组件中的eps值(默认=1e-5)。
batch_first: 如果设为`True`,则输入和输出张量将按照(batch, seq, feature)的形式提供。
默认:`False`。
Examples::
>>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
>>> src = torch.rand(10, 32, 512)
>>> out = encoder_layer(src)
Alternatively, when ``batch_first`` is ``True``:
>>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
>>> src = torch.rand(32, 10, 512)
>>> out = encoder_layer(src)
__init__:
self_attn: 使用 MultiheadAttention 实现的自注意力机制。
linear1, linear2: 用于前馈网络的线性层。
dropout: 用于前馈网络中的Dropout操作。
norm1, norm2: 层归一化层,用于自注意力和前馈网络之后。
dropout1, dropout2: 用于残差连接后的Dropout操作。
activation: 激活函数,根据传入的activation参数选择。
"""
__constants__ = ['batch_first']
def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu",
layer_norm_eps=1e-5, batch_first=False,
device=None, dtype=None) -> None:
factory_kwargs = {'device': device, 'dtype': dtype}
super(TransformerEncoderLayer, self).__init__()
self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
**factory_kwargs)
# Implementation of Feedforward model
self.linear1 = Linear(d_model, dim_feedforward, **factory_kwargs)
self.dropout = Dropout(dropout)
self.linear2 = Linear(dim_feedforward, d_model, **factory_kwargs)
self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.dropout1 = Dropout(dropout)
self.dropout2 = Dropout(dropout)
self.activation = _get_activation_fn(activation)
def __setstate__(self, state):
if 'activation' not in state:
state['activation'] = F.relu
super(TransformerEncoderLayer, self).__setstate__(state)
def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
r"""Pass the input through the encoder layer.
Args:
src: the sequence to the encoder layer (required).
src_mask: the mask for the src sequence (optional).
src_key_padding_mask: the mask for the src keys per batch (optional).
Shape:
see the docs in Transformer class.
forword flow:
首先通过自注意力层(self_attn)处理输入src,得到src2。
将src与经过Dropout操作的src2相加,然后通过层归一化(norm1)。
接着,src通过前馈网络,依次经过linear1、激活函数、dropout、linear2、dropout2,再与src相加,最后通过层归一化(norm2)。
"""
src2 = self.self_attn(src, src, src, attn_mask=src_mask,
key_padding_mask=src_key_padding_mask)[0]
src = src + self.dropout1(src2)
src = self.norm1(src)
src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
src = src + self.dropout2(src2)
src = self.norm2(src)
return src
class TransformerDecoderLayer(Module):
r"""`TransformerDecoderLayer` 由自注意力(self-attn)、多头注意力(multi-head-attn)和前馈网络(feedforward network)组成。
这个标准的解码器层基于论文 "Attention Is All You Need"。
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
Lukasz Kaiser, 和 Illia Polosukhin。2017年。注意力就是你所需要的一切。在《神经信息处理系统进展》期刊中,
第6000至6010页。用户在应用过程中可以对其进行修改或以不同方式实现。
参数:
d_model: 输入中预期特征的数量(必需)。
nhead: 多头注意力模型中的头数(必需)。
dim_feedforward: 前馈网络模型的维度(默认=2048)。
dropout: dropout的值(默认=0.1)。
activation: 中间层的激活函数,可选relu或gelu(默认=relu)。
layer_norm_eps: 层归一化组件中的eps值(默认=1e-5)。
batch_first: 如果设为`True`,则输入和输出张量将按照(batch, seq, feature)的形式提供。
默认:`False`。
Examples::
>>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
>>> memory = torch.rand(10, 32, 512)
>>> tgt = torch.rand(20, 32, 512)
>>> out = decoder_layer(tgt, memory)
Alternatively, when ``batch_first`` is ``True``:
>>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8, batch_first=True)
>>> memory = torch.rand(32, 10, 512)
>>> tgt = torch.rand(32, 20, 512)
>>> out = decoder_layer(tgt, memory)
__init__:
self_attn: 自注意力机制,用于处理目标序列内部的关系。
multihead_attn: 多头注意力机制,用于处理目标序列和记忆序列之间的关系。
linear1, linear2: 前馈网络的两个线性层。
dropout: 用于前馈网络中的Dropout操作。
norm1, norm2, norm3: 三个层归一化层,分别应用于自注意力、多头注意力和前馈网络之后。
dropout1, dropout2, dropout3: 用于残差连接后的Dropout操作。
activation: 激活函数,根据activation参数确定。
"""
__constants__ = ['batch_first']
def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu",
layer_norm_eps=1e-5, batch_first=False, device=None, dtype=None) -> None:
factory_kwargs = {'device': device, 'dtype': dtype}
super(TransformerDecoderLayer, self).__init__()
self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
**factory_kwargs)
self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
**factory_kwargs)
# Implementation of Feedforward model
self.linear1 = Linear(d_model, dim_feedforward, **factory_kwargs)
self.dropout = Dropout(dropout)
self.linear2 = Linear(dim_feedforward, d_model, **factory_kwargs)
self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.norm3 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.dropout1 = Dropout(dropout)
self.dropout2 = Dropout(dropout)
self.dropout3 = Dropout(dropout)
self.activation = _get_activation_fn(activation)
def __setstate__(self, state):
if 'activation' not in state:
state['activation'] = F.relu
super(TransformerDecoderLayer, self).__setstate__(state)
def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None,
tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
r"""使输入(及掩码)通过解码器层进行处理。
参数:
tgt: 需要送入解码器层的序列(必需)。
memory: 来自编码器最后一层的序列(必需)。
tgt_mask: 对于目标序列的掩码(可选)。
memory_mask: 对于记忆序列的掩码(可选)。
tgt_key_padding_mask: 每个批次的目标序列键的掩码(可选)。
memory_key_padding_mask: 每个批次的记忆序列键的掩码(可选)。
形状:
请参阅Transformer类中的文档。
流程:
首先,目标序列通过自注意力机制处理,结果与原目标序列残差连接后进行层归一化。
然后,处理后的目标序列与记忆序列通过多头注意力机制交互,结果与目标序列残差连接后进行层归一化。
最后,目标序列通过前馈神经网络,包括线性层、激活函数、Dropout和另一个线性层,处理后的结果与目标序列残差连接后进行层归一化。
"""
tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
key_padding_mask=tgt_key_padding_mask)[0]
tgt = tgt + self.dropout1(tgt2)
tgt = self.norm1(tgt)
tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
key_padding_mask=memory_key_padding_mask)[0]
tgt = tgt + self.dropout2(tgt2)
tgt = self.norm2(tgt)
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
tgt = tgt + self.dropout3(tgt2)
tgt = self.norm3(tgt)
return tgt
def _get_clones(module, N):
return ModuleList([copy.deepcopy(module) for i in range(N)])
def _get_activation_fn(activation):
if activation == "relu":
return F.relu
elif activation == "gelu":
return F.gelu
raise RuntimeError("activation should be relu/gelu, not {}".format(activation))