赞
踩
解析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))
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。