解析torch官方代码脚本文件:transformer.py。版本:1.9.1+cu111。
首先查看《Torch中多头注意力MultiheadAttention
的中文注释》解析;
最后查看下方transformer解析。
话不多说,看代码吧!
import copy
from typing import Optional, Anyimport 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 LayerNormclass 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 inhttps://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_encoderelse: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_decoderelse: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_modelself.nhead = nheadself.batch_first = batch_firstdef 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 outputdef 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 maskdef _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_layersself.norm = normdef 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 = srcfor 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 outputclass 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_layersself.norm = normdef 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 = tgtfor 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 outputclass 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 modelself.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.relusuper(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 srcclass 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 modelself.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.relusuper(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 tgtdef _get_clones(module, N):return ModuleList([copy.deepcopy(module) for i in range(N)])def _get_activation_fn(activation):if activation == "relu":return F.reluelif activation == "gelu":return F.geluraise RuntimeError("activation should be relu/gelu, not {}".format(activation))