positional encoding代码

以下是一个简单的 Python 代码，用于实现 Transformer 中的 Positional Encoding：

import torch
import math

class PositionalEncoding(torch.nn.Module):
    def __init__(self, d_model, max_seq_len=200):
        super().__init__()
        self.d_model = d_model
        self.max_seq_len = max_seq_len
        
        # Compute the positional encodings once 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)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        # Add positional encoding to all inputs
        x = x * math.sqrt(self.d_model)
        seq_len = x.size(1)
        if seq_len > self.max_seq_len:
            # Truncate long sequences
            x = x[:, :self.max_seq_len, :]
        else:
            # Pad short sequences
            padding = torch.zeros(x.size(0), self.max_seq_len - seq_len, x.size(2), device=x.device)
            x = torch.cat((x, padding), dim=1)
        x = x + self.pe[:seq_len, :]
        return x

这个代码中的 PositionalEncoding 类使用了 PyTorch 的 nn.Module 类，所以可以很容易地在 PyTorch 模型中使用。在构造函数中，我们首先计算了一个大小为 `(max_seq_len, d_model)` 的矩阵 `pe`，其中 `max_seq_len` 是序列的最大长度，`d_model` 是嵌入向量的维度。这个矩阵是通过计算一些正弦和余弦函数得到的，公式为：

$$
PE_{pos, 2i} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right) \\
PE_{pos, 2i+1} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)
$$

其中 $pos$ 是位置，$i$ 是维度。这个公式可以保证对于不同的位置和维度，得到的值是不同的，从而使得模型可以学习到位置信息。

在 `forward` 方法中，我们把输入张量 `x` 和 `pe` 相加，并返回结果。在相加之前，我们还对 `x` 进行了一些处理，包括将其乘以 $\sqrt{d_{model}}$（这个处理与 Transformer 中的 Multi-Head Attention 有关），以及根据 `max_seq_len` 对序列进行截断或填充。