transformer位置编码代码

以下是 Transformer 中位置编码的 Python 代码实现：

import torch
import math

class PositionalEncoding(torch.nn.Module):
    def __init__(self, d_model, max_seq_len):
        super().__init__()

        # Compute the positional encodings in log space
        pe = torch.zeros(max_seq_len, d_model)
        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)

        # Register the buffer so that it is not counted in the model parameters
        self.register_buffer('pe', pe)

    def forward(self, x):
        # Add the positional encodings to the input tensor
        x = x + self.pe[:x.size(0), :]
        return x

在这个实现中，我们首先计算了一个位置编码矩阵 `pe`，其中每一行代表一个位置的编码。然后，我们在 `forward` 函数中将这个矩阵加到输入张量 `x` 中。

这个实现中的位置编码采用了 Transformer 论文中的方式，即使用正弦和余弦函数来计算。具体来说，我们首先计算了一组 `div_term`，它们用于计算正弦和余弦函数中的除数，然后根据位置和这组除数计算出正弦和余弦的值，并将它们交替存储在 `pe` 矩阵中。最后，我们将 `pe` 矩阵加到输入张量中。