torch.nn.modules.transformer 的源代码

import copy
from typing import Optional, Any, Union, Callable

import torch
import warnings
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

__all__ = ['Transformer', 'TransformerEncoder', 'TransformerDecoder', 'TransformerEncoderLayer', 'TransformerDecoderLayer']

def _generate_square_subsequent_mask(
        sz: int,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
) -> Tensor:
    r"""生成一个用于序列的方形因果掩码。

    被掩码的位置填充为float('-inf')。未被掩码的位置填充为float(0.0)。
    """
    if device is None:
        device = torch.device('cpu')
    if dtype is None:
        dtype = torch.float32
    return torch.triu(
        torch.full((sz, sz), float('-inf'), dtype=dtype, device=device),
        diagonal=1,
    )


def _get_seq_len(
        src: Tensor,
        batch_first: bool
) -> Optional[int]:

    if src.is_nested:
        return None
    else:
        src_size = src.size()
        if len(src_size) == 2:
            # 未批处理: S, E
            return src_size[0]
        else:
            # 批处理: B, S, E if batch_first else S, B, E
            seq_len_pos = 1 if batch_first else 0
            return src_size[seq_len_pos]


class Transformer(Module):
    r"""一个transformer模型。

    用户可以根据需要修改属性。该架构基于论文 "Attention Is All You Need"。Ashish Vaswani, Noam Shazeer,
    Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, 和
    Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
    Processing Systems, pages 6000-6010.

    参数:
        d_model: 编码器/解码器输入的预期特征数量（默认=512）。
        nhead: 多头注意力模型中的头数（默认=8）。
        num_encoder_layers: 编码器中的子编码器层数（默认=6）。
        num_decoder_layers: 解码器中的子解码器层数（默认=6）。
        dim_feedforward: 前馈网络模型的维度（默认=2048）。
        dropout: dropout值（默认=0.1）。
        activation: 编码器/解码器中间层的激活函数，可以是字符串
            （"relu" 或 "gelu"）或一元可调用对象。默认: relu
        custom_encoder: 自定义编码器（默认=None）。
        custom_decoder: 自定义解码器（默认=None）。
        layer_norm_eps: 层归一化组件中的eps值（默认=1e-5）。
        batch_first: 如果为 ``True``，则输入和输出张量以 (batch, seq, feature) 提供。默认: ``False`` (seq, batch, feature)。
        norm_first: 如果为 ``True``，编码器和解码器层将在其他注意力和前馈操作之前执行LayerNorms，否则在之后。默认: ``False`` (after)。
        bias: 如果设置为 ``False``，``Linear`` 和 ``LayerNorm`` 层将不会学习加性偏置。默认: ``True``。

    示例::
        >>> 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)

    注意: 一个完整的示例，使用 nn.Transformer 模块进行词语言模型的应用，可以在
    https://github.com/pytorch/examples/tree/master/word_language_model
    """

    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: Union[str, Callable[[Tensor], Tensor]] = F.relu,
                 custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None,
                 layer_norm_eps: float = 1e-5, batch_first: bool = False, norm_first: bool = False,
                 bias: bool =</span