Transformer 学习

import torch
sentence = 'Life is short, eat dessert first'
sentence_split = sentence.replace(',', '').split()
dc = {s:i for i,s in enumerate(sorted(sentence_split))}

sentence_int = torch.tensor([dc[s] for s in sentence_split])
print(sentence_int)

torch.manual_seed(123)
embed = torch.nn.Embedding(6, 16)
embedded_sentence = embed(sentence_int).detach()

import torch
sentence = 'Life is short, eat dessert first'
sentence_split = sentence.replace(',', '').split()
# 使用当前sentence建立word-to-pos dictionary
dc = {s:i for i,s in enumerate(sorted(sentence_split))}
# 使用当前dictionary中每个单词的位置
sentence_int = torch.tensor([dc[s] for s in sentence_split])
print(sentence_int)

torch.manual_seed(123)
# Embedding会将word index映射到vector
embed = torch.nn.Embedding(len(sentence_split), 16)
embedded_sentence = embed(sentence_int).detach() # 6x16 array

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

class PositionalEncoding(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        # 注意这里只计算偶数位置，对应就是公式中的2i。
        # 另外这里用了exp(i*log10000)=pow(10000,i)来计算后者
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
        )
        # 按照位置的奇偶性来分配sin or cos计算
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0) # 最后增加最前面维度，方便和输入的x相加
        self.register_buffer("pe", pe)

    def forward(self, x):
        # x: (b,T,d_model), pe: (1,max_len,d_model)
        # T通常是这个batch中的所有sentences的max length（其实经常它就等于max_len）
        # max_len则通常是整个数据集全部sentences的max length，所以max_len>=T
        # 因此下面两者可以相加
        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return self.dropout(x)

# Attention定义
def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    # Q,K,V的维度个数不确定，self attention时候是(b,n,d)，而multi-head attention
    # 则会是 (b,h,n,d)。不过无论哪一种，我们只针对最后两个维度做矩阵相乘，因此只需要转置
    # K的最后两维。而 torch.matmul(A,B) 计算中，如果A,B前1个或2个维度相同，那么它只做
    # batch multiplication，即只针对每个batch的最后两个维度做乘法，保留前面的维度. 
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # mask在后面 Decoder中的attention才会用到，目的是消除和未来单词之间的关联，因为在
    # Decoder中，单词是逐个预测的，而不是像Encoder这样把整个句子全部单词全扔进去训练。
    # 由于scores含义就是每个单词和其他单词之间的关联程度，因此显然单词不能和未来单词相关联，
    # 于是要用mask（通常是一个下三角数组）将上三角部分给清掉。这里用了-1e9一个很大的负数，是为了
    # 使得计算 softmwax 后这个值变成0.
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    # softmax针对最后一个维度
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    # (n,n) x (n,d_v) = (n,d_v)
    return torch.matmul(p_attn, value), p_attn

# Multi-Head Attention中生成Q,K,V并调用Attention
class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        ...
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        ...

    def forward(self, query, key, value, mask=None):
        # 通常输入的Q,K,V还是三个维度 (batch_size, seq_len, d_model)
        if maskisnotNone:
            # mask: (batch_size, 1, seq_len, seq_len), for h heads
            mask= mask.unsqueeze(1)
        batch_size= query.size(0)

        # 对Q,K,V都经过对应的Linear层。下面是一个简洁写法。
        # (batch_size, seq_len, d_model) -> (batch_size, seq_len, h, d_k)
        #NOTE: 1) these two lines are equivalent to the following commented lines for each of query, key, value:
        # 2) the zip(A,B) operation can align the length of A and B by taking the minimum length of the two, and
        #   here A,B could be either list or tuples.
        query, key, value= [
            lin(x).view(batch_size,-1,self.h,self.d_k).transpose(1, 2) for lin, xin zip(self.linears, (query, key, value))
        ]
        """
        # 上面的简单写法展开后如下所示，下面是以Query为例，K,V类似。
        ## Split the above lines into the following commented lines for each of query, key, value:
        # (batch_size, seq_len, d_model)
        query = self.linears[0](query)
        # (batch_size, seq_len, h, d_k), since d_k = d_model // h
        query = query.view(batch_size, -1, self.h, self.d_k)
        # (batch_size, h, seq_len, d_k), since we want to compute attention on the last two dimensions, for each head
        query = query.transpose(1, 2)
        """

        # 直接调用attention()方法，里面已经设计好，可以处理多头机制
        # x: (batch_size, h, seq_len, d_k), attn: (batch_size, h, seq_len, seq_len)
        x,self.attn= attention(query, key, value, mask=mask, dropout=self.dropout)

        # 合并全部heads的结果，然后恢复 (batch_size, seq_len, d_model) 的形式
        # concat results of all heads together and merge the two dimensions.
        # (batch_size, h, seq_len, d_k) -> (batch_size, seq_len, d_model), since d_model = h * d_k
        # What this line does is:
        # 1) transpose(1, 2): Swaps dimensions 1 and 2, changing tensor from shape [batch_size, h, seq_len, d_k] to [batch_size, seq_len, h, d_k]
        # 2) contiguous(): Makes the tensor contiguous in memory after transposing, which is a MUST and required for view() to work correctly, otherwise
        #   it will raise an error during runtime (since transpose() only changes the "viewing" shape of the tensor but not the memory layout, while the
        #   view() function MUST take a continuous tensor to reshape).
        # 3) view(batch_size, -1, self.h * self.d_k): Reshapes the tensor to [batch_size, seq_len, d_model] where d_model = h * d_k
        #    The -1 automatically calculates the seq_len dimension based on the total number of elements
        x= x.transpose(1, 2).contiguous().view(batch_size,-1,self.h*self.d_k)

        # It might be not that necessary to release memory here, but just to be safe that they won't be used anywhere else after this.
				del query, key, value

        # 最后经过Linear层
        # Remember to run the last linear layer. (batch_size, seq_len, d_model) -> (batch_size, seq_len, d_model)
				return self.linear[-1](x)

# SublayerConnection用于封装几个相关操作：dropout, LayerNorm以及Residual Connection
class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        # 注意LayerNorm只作用在最后一个维度即 d_model 维度上（实现中size=d_model)
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        # 这里定义了LayerNorm以及Add Residual的重复操作，并且输入sublayer，它可以
        # 是Multi-Headed Attention，也可以是 Feed Forward，这样就避免重复实现，使代码简洁
        return x + self.dropout(sublayer(self.norm(x)))

# FeedForward实现
class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."

    def __init__(self, d_model, d_ff, dropout=0.1):
        # 通常中间会升维，例如d_model=512, d_ff=2048
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(self.w_1(x).relu()))

# 调用就很方便了
...
self.sublayers = Clones(SublayerConnection(size, dropout), 2)
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # Multi-Headed attention 
x = self.sublayer[1](x, self.feed_forward) # Feed Forward

def clones(module, N):
    """
    Produce N identical layers.
    复制N个module模块返回list of modules. 用处很多所以为了简洁定义了这个函数
    """
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

def attention(query, key, value, mask=None, dropout=None):
    """

    Args:
        query: size [b, n, d_model] for single head, or [b, h, n, d_q] for multi-headed
        key: size [b, m, d_model] for single head, or [b, h, m, d_k] for multi-headed. d_q = d_k
        value: size [b, m, d_model] for single head or [b, h, m, d_v] for multi-headed.
        mask: size [b, 1, m] for src_mask, or [b, n, n] for tgt_mask
        dropout: nn.Dropout, dropout layer

        NOTE: Transformer 流程中有 2 Self Attentions 和 1 Cross Attention:
        self-attention: Q,K,V 来源相同(都是src 或都是 tgt), 此时m=n, 且 mask 会对应 src_mask or tgt_mask
        cross attention: 那么通常 Q 来自 tgt 而 K,V 来自 src, 此时 n=tgt_seq_len, m=src_seq_len, 两者通常不同,
        而对应的 mask 是 src_mask

        实现中的 d_q, d_k, d_v 通常是完全相等的, 且 d_q = d_k = d_v = d_model // h, h is the number of heads
        例如 d_model = 512, h = 8, d_q = d_k = d_v = 64

    Returns:
        output: torch.Tensor, attention output, [b, n, d_model] or [b, h, n, d_v]
        p_attn: torch.Tensor, attention weights (scores), size [b, n, m] or [b, h, n, m]
    """
    d_k = query.size(-1)

    # Q: (b,h,n,d_q), K: (b,h,m,d_k), d_q = d_k, 而 n,m 在 cross attention 时不等
    # NOTE: matmul(A,B)中, 如果A,B前面1-2维度完全相同并且后两维可以支持数组相乘, 那么它的结果就是 multiplication per batch,
    # 即它会认为前面两维 b,h 是 batch 然后对后两维相乘. 结果就是 (b,h,n,m)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        # NOTE: 两种情况
        # 1) self attention: mask=src_mask or tgt_mask: [b, 1, n] or [b, n, n], 此时 scores都是[b, h, n, n] (n=m)
        # 2) cross attention: mask=src_mask: [b,1,m], 此时 scores=[b, h, n, m]
        # "masked_fill()" supports broadcasting, 所以无论哪种情况都可以运行
        #
        # 注意这里将等于0的位置设置为-inf即负最小值，而不是0。这是因为，此后scores要用来做softmax，其计算公式是 e^x/sum(e^x)，而
        # 我们期待是经过softmax后将这些invalid的位置设为0，而e^(-inf)=0，故这里设置为-inf
        scores = scores.masked_fill(mask == 0, -1e9)

    # 只针对最后一个维度计算 softmax
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)

    # 最后和V相乘, [b,h,n,m] * [b,h,m,d_v] -> [b,h,n,d_v]
    return torch.matmul(p_attn, value), p_attn

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        # h: number of heads (e.g. 8), d_model: embedding size (e.g. 512), dropout: dropout ratio (e.g. 0.1)
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0

        # We assume d_v always equals d_k
        self.d_k = d_model // h  # e.g. 512 // 8 = 64
        self.h = h  # e.g. 8

        # Clone a list of 4 Linear layers. First 3 for Q,K,V, last for the final linear output layer
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        """
        Args:
            query: torch.Tensor, query tensor, size [batch_size, seq_len, d_model]
            key: torch.Tensor, key tensor, size [batch_size, seq_len, d_model]
            value: torch.Tensor, value tensor, size [batch_size, seq_len, d_model]
            mask: torch.Tensor, mask tensor to mask out certain positions, size [batch_size, seq_len, seq_len]

        Returns:
            output: torch.Tensor, attention output, size [batch_size, seq_len, d_model]
        """
        if mask is not None:
            # mask: (batch_size, 1, seq_len, seq_len), for h heads
            mask = mask.unsqueeze(1)
        batch_size = query.size(0)

        # (batch_size, seq_len, d_model) -> (batch_size, seq_len, h, d_k)
        # NOTE: 1) these two lines are equivalent to the following commented lines for each of query, key, value:
        # 2) the zip(A,B) operation can align the length of A and B by taking the minimum length of the two, and
        #   here A,B could be either list or tuples.
        query, key, value = [
            lin(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (query, key, value))
        ]
        """
        ## Split the above lines into the following commented lines for each of query, key, value:
        # (batch_size, seq_len, d_model)
        query = self.linears[0](query)
        # (batch_size, seq_len, h, d_k), since d_k = d_model // h
        query = query.view(batch_size, -1, self.h, self.d_k)
        # (batch_size, h, seq_len, d_k), since we want to compute attention on the last two dimensions, for each head
        query = query.transpose(1, 2)
        """

        # x: (batch_size, h, seq_len, d_k), attn: (batch_size, h, seq_len, seq_len)
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)

        # concat results of all heads together and merge the two dimensions.
        # (batch_size, h, seq_len, d_k) -> (batch_size, seq_len, d_model), since d_model = h * d_k
        # What this line does is:
        # 1) transpose(1, 2): Swaps dimensions 1 and 2, changing tensor from shape [batch_size, h, seq_len, d_k] to [batch_size, seq_len, h, d_k]
        # 2) contiguous(): Makes the tensor contiguous in memory after transposing, which is a MUST and required for view() to work correctly, 
        # otherwise it will raise an error during runtime (since transpose() only changes the "viewing" shape of the tensor but not the memory 
        # layout, while the view() function MUST take a continuous tensor to reshape).
        # 3) view(batch_size, -1, self.h * self.d_k): Reshapes the tensor to [batch_size, seq_len, d_model] where d_model = h * d_k
        #    The -1 automatically calculates the seq_len dimension based on the total number of elements
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)

        # It might be not that necessary to release memory here, but just to be safe that they won't be used anywhere else after this.
        del query, key, value

        # Remember to run the last linear layer. (batch_size, seq_len, d_model) -> (batch_size, seq_len, d_model)
        return self.linears[-1](x)

# SublayerConnection封装几个操作：dropout, LayerNorm以及Residual Connection
class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))

# FeedForward实现
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        # 通常中间会升维，例如d_model=512, d_ff=2048
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(self.w_1(x).relu()))

class EncoderLayer(nn.Module):
    """
    包括两个步骤:
    1) MultiHeadedAttention -> SublayerConnection (包含 LayerNorm -> Add Residual)
    2) FeedForward -> SublayerConnection (包含 LayerNorm -> Add Residual)
    Dropout 层已经在 SublayerConnection 添加了, 这里只是传入 dropout ratio
    """

    def __init__(self, size, self_attn, feed_forward, dropout):
        # size: embedding size, 通常等于 d_model, self_attn 和 feed_forward 要提前定义好, dropout: dropout ratio
        # 这里的 size 只用来传入 Encoder 中用于其 LayerNorm 的定义, 而 EncoderLayer 中则不用
        super(EncoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)

    def forward(self, x, mask):
        # x: (batch_size, seq_len, d_model)
        # mask: (batch_size, 1, seq_len) or (batch_size, seq_len, seq_len)
        # return: (batch_size, seq_len, d_model)
        #! NOTE: 注意 SublayerConnection() 的第二个参数是关于 attention 的函数而不是数值。
        #! 这是因为，从SublayerConnection.forward()函数不难看出，它的第二个参数sublayer在里面用法是
        #! sublayer(.)，即它应该是一个 Layer module, 因此设计成函数作为输入便可以用做 Layer module
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        # 第2步就是 FeedForward -> SublayerConnection
        return self.sublayer[1](x, self.feed_forward)

class Encoder(nn.Module):
    """
    这就是 Encoder 最顶层的类了, 它包含 N 个相同的 EncoderLayer, 例如 Transformer 中一个例子 N=6
    """

    def __init__(self, layer, N):
        # 调用时, layer 通常是提前定义好的 EncoderLayer, N的一个例子 N=6
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        # x: (batch_size, seq_len, d_model)
        # mask: (batch_size, 1, seq_len) or (batch_size, seq_len, seq_len), depends on self.layers API
        # return: (batch_size, seq_len, d_model)
        # 输入的x,mask这个接口完全取决于EncoderLayer.forward()的定义
        for layer in self.layers:
            x = layer(x, mask)

        # 注意整个 Encoder 结束后, 官方代码又加上了一层 Layer Norm, 不过流程图上并没有。或许是考虑到
        # 接下来 Encoder 的输出会被传入到 Decoder Attention 作为输入, 因此加上 LayerNorm 使输出更稳定
        return self.norm(x)

def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):
    """
    Args:
        src_vocab: source vocabulary size,
        tgt_vocab: target vocabulary size
        N: number of repeated EncoderLayer and DecoderLayer modules,
        d_model: embedding size
        d_ff: feed forward layer size,
        h: number of heads
        dropout: dropout ratio
    """
    c = copy.deepcopy
    **attn = MultiHeadedAttention(h, d_model)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    pos_embed = PositionalEncoding(d_model, dropout)**

    **encoder = Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N)**
    decoder = Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N)
    src_embed = nn.Sequential(Embeddings(src_vocab, d_model), pos_embed)
    tgt_embed = nn.Sequential(Embeddings(tgt_vocab, d_model), pos_embed)
    generator = Generator(d_model, tgt_vocab)

    model = Transformer(encoder, decoder, src_embed, tgt_embed, generator)

    # This was important from their code.
    # Initialize parameters with Glorot / fan_avg.
    # Paper title: Understanding the difficulty of training deep feedforward neural networks Xavier
    # 官方代码中的一个初始化方法, 还留在这不动
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model

#Encoder 中的src_mask的设计：很简单，只需要屏蔽invalid token即id=0的地方即可，这里通常pad=0
src_mask= (src!= pad).unsqueeze(-2)

# Decoder tgt_mask 设计
def make_std_mask(tgt: torch.Tensor, pad: int) -> torch.Tensor:
    # [batch_size, 1, seq_len], all original non-zero elements are 1, and the rest 0 (by default pad=0)
    # NOTE 本方法其实只是为了建立 Decoder 中使用的 tgt_mask，而它既要屏蔽invalid tokens（包括 invalid word 以及 pad，
    # 而这两者对应的 id 都是 0），还要屏蔽未来的words。下面这行就是为了前者。
    tgt_mask = (tgt != pad).unsqueeze(-2) # (b, 1, seq_len)

    # NOTE: here the subsequent_mask() returns the look-ahead mask [1, seq_len, seq_len], with
    # all lower left triangle is 1 (include diagonal) the rest 0. So how and why this "&" operation works
    # on two tensors with size [batch_size, 1, seq_len] and [1, seq_len, seq_len]? The answer is that,
    # pytorch will broadcast the two tensors to fit each other's size. That is, it first broadcasts the left
    # tensor [batch_size, 1, seq_len] to [batch_size, seq_len, seq_len], and also broadcast the right tensor
    # [1, seq_len, seq_len] to [batch_size, seq_len, seq_len], and then do the "&" operation. So the final result
    # is a [batch_size, seq_len, seq_len].
    # 这里 & 后面的mask就是为了屏蔽未来words，因此设置成了一个binary下三角阵，下三角部分和对角线上都是1，其余是0。两个 masks做了
    # 且运算后得到最终的mask，兼顾了两种功能。
    tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)
    return tgt_mask # (b, seq_len, seq_len)

def subsequent_mask(size):
    # 计算binary下三角阵，下三角和对角线是1，其余是0
    attn_shape = (1, size, size)
    # triu() 生成上三角矩阵，diagnoal=1表明去掉对角线，其结果是只有右上三角是1（不包括对角线），其余是0。
    # 接下来返回的是 subsequent_mask == 0，因此取反后，只有左下角和对角线是1，其余是0。
    subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
        torch.uint8
    )
    return subsequent_mask == 0

class DecoderLayer(nn.Module):
    """
    和 EncoderLayer 类似但是会多一个步骤, 包括3个步骤:
    1) MultiHeadedAttention (self attention with target data only) -> SublayerConnection (包含 LayerNorm -> Add Residual)
    2) MultiHeadedAttention (cross attention, takes input both Encoder output and the last self attention output) -> SublayerConnection
    3) FeedForward -> SublayerConnection

    Dropout 层已经在 SublayerConnection 添加了, 这里只是传入 dropout ratio
    """

    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        # size: embedding size, 通常等于 d_model
        # self_attn, src_attn, feed_forward 要提前定义好, 其中 src_attn 是使用了 Encoder source data 的 attention 机制
        # dropout: dropout ratio
        # 类似的, 这里的 size 只用来传入 Decoder 中用于其 LayerNorm 的定义, 而 DecoderLayer 中则不用
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, src_mask, tgt_mask):
        # x: (batch_size, seq_len, d_model)
        # memory: (batch_size, seq_len, d_model)
        # src_mask: (batch_size, 1, seq_len, seq_len), tgt_mask: (batch_size, seq_len, seq_len)
        # return: (batch_size, seq_len, d_model)
        #! NOTE: 同样类似 Encoder 中的定义, 这里的某些 SublayerConnection() 的第二个参数是关于attention的函数而不是数值。
        #! 这是为了将其作为类似一个 Layer 的形式传入

        # 1) MultiHeadedAttention (self attention) -> SublayerConnection (包含 LayerNorm -> Add Residual)
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        # 2) MultiHeadedAttention (cross attention) -> SublayerConnection
        x = self.sublayer[1](x, lambda x: self.src_attn(x, memory, memory, src_mask))
        # 3) FeedForward -> SublayerConnection
        return self.sublayer[2](x, self.feed_forward)

class Decoder(nn.Module):
    """
    这就是 Decoder 最顶层的类了, 和 Encoder 很类似, 它包含 N 个相同的 DecoderLayer, 例如 N=6
    """

    def __init__(self, layer, N):
        # layer 通常是提前定义好的 DecoderLayer, N的一个例子 N=6
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, src_mask, tgt_mask):
        # x: (batch_size, seq_len, d_model), memory: (batch_size, seq_len, d_model)
        # src_mask: (batch_size, seq_len, seq_len), tgt_mask: (batch_size, seq_len, seq_len)
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)

        # 类似 Encoder, 这里当 Decoder 结束后, 官方代码也给加上了一层 Layer Norm
        return self.norm(x)

class Generator(nn.Module):
    """
    Generator 就是 Decoder 之后到最终输出的这几层, 很简单只包含 Linear -> Softmax. 其中 Linear 作用是将
    最后一维的 Embedding size (d_model) 重新变为 vocab_size, 而 Softmax 则是将数值重新变成概率值
    """

    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        # return: (batch_size, seq_len, vocab)
        # 注意最后使用了 log(softmax) 而不是 softmax
        return F.log_softmax(self.proj(x), dim=-1)

class Transformer(nn.Module):
    """
    Transformer 就是 Encoder + Decoder + Generator. 官方代码把这个类定义为 EncoderDecoder, 不过感觉还是直接 Transformer 更好些
    另外参考的 English-Chinese Translator 代码也是如此定义的
    """

    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        # encoder, decoder, src_embed, tgt_embed, generator 都是提前定义好的. 注意:
        # 1) src_embed, tgt_embed 都包括了 Word Embedding -> Positional Embedding 这两个步骤,
        # 通常可以用 src_embed = nn.Sequential(Embedding, PositionalEmbedding) 这样提前定义好然后传入
        # 2) generator只是传入进来但没有包含到流程中, 它会在外面显式调用, 方便操作

        super(Transformer, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def encode(self, src, src_mask):
        # src: (batch_size, seq_len)
        # src_mask: (batch_size, 1, seq_len) or (batch_size, seq_len, seq_len)
        # return: (batch_size, seq_len, d_model)
        # NOTE 注意 src 是 raw word ID 数组, 因此预先要做 Embedding
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        # memory: (batch_size, seq_len, d_model), basically the Encoder result on src
        # src_mask: (batch_size, 1, seq_len) or (batch_size, seq_len, seq_len)
        # tgt: (batch_size, seq_len)
        # tgt_mask: (batch_size, seq_len, seq_len)
        # return: (batch_size, seq_len, d_model)
        # NOTE 注意 tgt 也是 raw word ID 数组, 因此预先要做 Embedding
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

    def forward(self, src, tgt, src_mask, tgt_mask):
        # src: (batch_size, seq_len), tgt: (batch_size, seq_len)
        # src_mask: (batch_size, seq_len, seq_len), tgt_mask: (batch_size, seq_len, seq_len)
        # return: (batch_size, seq_len, d_model)
        # 注意它只有 Encode -> Decoder 流程, 不包括 Generator, 后者会在外面显式调用
        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)

def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):
    """
    Args:
        src_vocab: source vocabulary size,
        tgt_vocab: target vocabulary size
        N: number of repeated EncoderLayer and DecoderLayer modules,
        d_model: embedding size
        d_ff: feed forward layer size,
        h: number of heads
        dropout: dropout ratio
    """
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    pos_embed = PositionalEncoding(d_model, dropout)

    encoder = Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N)
    decoder = Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N)
    src_embed = nn.Sequential(Embeddings(src_vocab, d_model), pos_embed)
    tgt_embed = nn.Sequential(Embeddings(tgt_vocab, d_model), pos_embed)
    generator = Generator(d_model, tgt_vocab)

    model = Transformer(encoder, decoder, src_embed, tgt_embed, generator)

    # This was important from their code.
    # Initialize parameters with Glorot / fan_avg.
    # Paper title: Understanding the difficulty of training deep feedforward neural networks Xavier
    # 官方代码中的一个初始化方法, 还留在这不动
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model

class LabelSmoothing(nn.Module):
    """
    Implement label smoothing regularization technique.
    实现一个 smoothing distribution 来替代传统的 distribution 的比较方式. 具体见 forward() 中的记录
    """

	def __init__(self, size, padding_idx, smoothing=0.0):
        # size: vocabulary size 注意这个不是 d_model 了
        # padding_idx: index of padding token in vocabulary, usually 0
        # smoothing: smoothing parameter (epsilon), default is 0.0 (no smoothing) 通常用较小的值例如 0.1
        super(LabelSmoothing,self).__init__()

        # KLDivLoss expects log probabilities as input and probabilities as target
        # reduction="sum" means the loss will be summed over all elements
				self.criterion= nn.KLDivLoss(reduction="sum")
				self.padding_idx= padding_idx
				self.confidence= 1.0- smoothing  # Probability assigned to correct class
				self.smoothing= smoothing  # Smoothing parameter (epsilon)
				self.size= size  # Vocabulary size
				self.true_dist=None  # Store distribution for visualization/debugging

	def forward(self, x, target):
        # x: model output logits after log_softmax, shape [b, v], where v is vocabulary size
        # target: groundtruth true labels, shape [b]
        # return: loss value
        #NOTE: 通常 x 传入时就已经被reshape成了只有2个维度, 即类似 [b,n,v] -> [b*n,v], 故 x=[b,v] 中
        # 的 b 其实指的是 batch size * seq_len. target同理, 被reshape成了只有1个维度, 类似 [b,n] -> [b*n]

        # Ensure input has correct vocabulary size dimension
				assert x.size(1)==self.size

        # Create a copy of input tensor to build our smoothed target distribution
        # [b, v]
        true_dist= x.data.clone()

        # Fill with small probability for all classes
        # We divide by (size-2) to account for correct class and padding token
        # 默认全部赋予一个非常小的值, 例如 smoothing=0.1, size=10000, 那么这个值接近 1e-5
        # [b, v]
        true_dist.fill_(self.smoothing/ (self.size- 2))

        # Put the confidence (1-smoothing) probability on the correct class
        # unsqueeze(1) converts target from [batch_size] to [batch_size, 1]
        #!NOTE: self.scatter_(dim, index, src) is a very useful function, it will copy
        #! the value of src to the index position of the dim dimension.
        #! 它的作用是将 src 中的每个数值 src[i,j] 赋给 self 数组中的第 dim 维度下的第 index[i,j]
        #! 这个位置. 即: 它只修改第 dim 维度. 例如这里, true_dist is (b, v), dim=1, index is (b,1)
        #! 而 src=self.confidence 是一个 scalar 固定值. 因此, dim=1 说明修改的是 true_dist 的 v 这个
        #! 维度即 vocab. 具体的话就是 true_dist[0,index[0][0]] = src, true_dist[1,index[1][0]] = src, ...
        #! 而 index=target.data.unsqueeze(1) 其实就是 target word ID index.
        #! 总结的话就是, 把 true_dist 的维度 dim=1 这里维度上的, index = groundtruth word ID from target 的这些
        #! 位置上都赋值为 confidence 值, 例如 0.9 (剩下全部位置还是上面计算得到的一个极小值)
        true_dist.scatter_(1, target.data.unsqueeze(1),self.confidence)

        # Set probability for padding token to 0 将pad所在的位置全部设为0
        true_dist[:,self.padding_idx]= 0

        # Find positions where target is padding token
        # 找到 target 中的 pad 所在的全部位置的 indices
        mask= torch.nonzero(target.data==self.padding_idx)

        # If padding tokens exist in target
				if mask.dim()> 0:
            # index_fill_(dim, index, value) 是将第dim维度中的index数组记录的位置都设置为value
            # 这里就是将 true_dist 中的这些 (对应target中的) pad indices 都设为 0
            true_dist.index_fill_(0, mask.squeeze(), 0.0)

				self.true_dist= true_dist  # 这就是新的 Smoothed Groundtruth distribution

        # 计算 x 和新的 smoothed groundtruth distribution 的 KL Divergence Loss
        #!NOTE: 本函数的输入x通常是Transformer整个流程的输出结果, 其最后一层是 Generator. 从它的 forward 函数定义
        #! 看, 它的最后一层是 log_softmax, 即计算了 softmax 后再取了log. 即, x 的值是已经计算了log的. 不过，显然这里的
        #! true_dist并没有计算log, 那么下面这个方法为何正确呢？这是因为, criterion 即 nn.KLDivLoss(input, target) 比较特别
        #! 它的输入 input 已经被要求要在 log-space 了, 而 target 则默认没有被要求(有个参数可以切换 target 是否已经
        #! 在这里面了）
				return self.criterion(x, true_dist)

class SimpleLossCompute:
    """
    A simple loss compute and train function based on the output of the model.

    NOTE 把计算 loss 的 Criterion 给封装到一个 wrapper 中, 好处是:
    1) 区分 training loss 和 evaluation loss 时的不同: 前者有 optimizer, 后者不需要.
    而其他流程都相同. 故封装成 wrapper 后定义两个 instances 来使用, 使得代码简洁
    2) 计算 loss 时候初始化的 generator, criterion, optimizer 都是固定值, 然后每次计算
    新的输入 x,y 之间的 loss, 因此用一个wrapper类先初始化这些固定参数, 然后调用它来计算新的
    (x,y) 的 loss 比较简单
    """

    def __init__(self, generator, criterion, opt=None):
        # generator: Generator module that projects to vocabulary size
        # criterion: Loss function (typically LabelSmoothing)
        # opt: Optimizer (optional, can be None for evaluation)
        self.generator = generator
        self.criterion = criterion
        self.opt = opt

    def __call__(self, x, y, norm_factor):
        # x: Transformer model output tensor (batch_size, seq_len, d_model)
        # y: target tensor with ground truth indices (batch_size, seq_len)
        # norm_factor: normalization factor, typically the number of tokens in the batch
        # return: loss value, one scalar
        #! NOTE: 注意这个函数被设计为 __call__() 作用是可以将这个类的 instance 用作类似一个函数,
        #! 其实就是类似 nn.Module 中的 forward() 函数, 只是这个类不是 nn.Module, 所以只能设计成
        #! __call__()函数来代替. 用法也一样, 例如: 定义 sc = SimpleLossCompute(...), 然后调用
        #! loss = sc(x,y,norm_factor) 就行了

        # Apply final projection to vocabulary size
        #! NOTE: 注意 Transformer 中的最后的 Generator 是在这里被调用的, 说实话感觉有些过于隐蔽
        #! 并且结构有些奇怪, 我还是感觉放在 Transformer 类里面的 forward() 更合适.
        x = self.generator(x)  # (batch_size, seq_len, vocab_size)

        # Reshape tensors for loss computation:
        # - Flatten batch and sequence dimensions to compute loss over all tokens
        # - contiguous() ensures memory layout is optimized for the view operation
        # NOTE: 计算 backward loss 时候通常要除以一个 norm_factor, 通常它是 #tokens in batch
        # 即计算的是 loss per token, 这是为了 stable gradient descent, 避免某些 batch 中的
        # tokens 过多导致 gradients 过大, 使得其他 tokens 的 gradients 被稀释.
        loss = (
            # Normalize loss by token count for stable gradients during training
            self.criterion(x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1))
            / norm_factor
        )

        # Update weights if optimizer is provided (training mode)
        # NOTE: training mode 才有 opt 传入, 而 validation 的话就不用
        if self.opt is not None:
            # Compute gradients through backpropagation only if optimizer is provided (training mode)
            loss.backward()

            self.opt.step()  # Apply gradients to update weights
            self.opt.optimizer.zero_grad()  # Reset gradients for next batch

        # Return the un-normalized loss value (multiplied back by norm_factor)
        # This gives the actual total loss for reporting/monitoring purposes
        # while keeping gradients normalized during the backward pass
        # NOTE: 最终返回的 loss 还会乘以 norm_factor, 使得最终的 loss 变回和 batch size 相关,
        # 这是因为这个最终的 loss 通常会被用来统计全部 batches 的 loss 总和, 因此要用原始的 loss
        return loss.item() * norm_factor

# 提前定义 Loss Instance
loss_compute = SimpleLossCompute(self.model.generator, criterion, opt=self.optimizer)

def run_epoch(self, data, loss_compute):
    # data: data iterator, yield batches of data
    # loss_compute: loss function and optimizer wrapper
    # epoch: current epoch number
    # return: loss per token
    start = time.time()
    total_loss = 0
    tokens = 0
    total_tokens = 0
    batch_interval = 50
    for i, batch in enumerate(data):
        batch = batch.to(self.device)
        out = self.model(batch.src, batch.tgt, batch.src_mask, batch.tgt_mask)
        # 调用时，follow __call__() API in SimpleLossCompute()，而 tgt_y 就是提到的每个位置错开一位的 tgt[:, 1:t] 这样
        loss = loss_compute(out, batch.tgt_y, batch.ntokens)
        total_loss += loss
        total_tokens += batch.ntokens
        tokens += batch.ntokens
        if i % batch_interval == 1:
            elapsed = time.time() - start
            print(
                f"  Epoch {self.current_epoch}, Batch {i}, tokens {batch.ntokens}, tokens/sec {tokens / elapsed:.3f}, loss {loss / batch.ntokens:.3f}"
            )
            start = time.time()
            tokens = 0

    # 最终返回 loss per token
  return total_loss / total_tokens

def evaluate(self, src, pad=0, bos_id=2, max_len=50):
    # src: input sentence word ID tensor, (b, src_len)
    # max_len: maximum output length
    # return: output sentence IDs, (b, tgt_len)
    self.model.eval()
    with torch.no_grad():
        src = torch.from_numpy(src).to(self.device)
        src_mask = (src != pad).unsqueeze(-2).to(src.device)  # (b, 1, src_len)
        memory = self.model.encode(src, src_mask)  # (b, src_len, d_model)
        # (b, 1), tgt data starts with bos_id ("BOS")
        tgt = (
            torch.ones(src.shape[0], 1)
            .fill_(bos_id)
            .type_as(src.data)
            .to(src.device)
        )
        for t in range(max_len - 1):
            tgt_mask = (
                subsequent_mask(tgt.size(-1)).type_as(src.data).to(src.device)
            )

            # (b, tgt_len, d_model)
            out = self.model.decode(memory, src_mask, tgt, tgt_mask)

            # NOTE 根据我们的 decoder 设计, 输入 tgt 中的 dim=1 即 tgt_len 这个维度
            # 的范围是 [0:t-1], 预测的 out 中的 tgt_len 这个维度的范围是 [1:t] (参见 Batch 中的 tgt 和 tgt_y)
            # 即我们用从 0 开始到当前的 t-1 时刻的所有单词来预测未来时刻 t 的单词. 因此, 每次预测后
            # 只需要位置 t 的这个新的预测值, 它正是 dim=1 这个维度的最后一个位置, 即 out[:, -1]
            # (b, d_model) -> (b, vocab_size)
            prob = self.model.generator(out[:, -1])

            # (b), positions of max elements along dim=1
            _, next_word = torch.max(prob, dim=1)

            # 将新的预测值合并历史记录, 构成新的 concat (b, tgt_len) with (b, 1) to get (b, tgt_len+1)
            tgt = torch.cat([tgt, next_word.unsqueeze(-1)], dim=1)

    return tgt