Learning PyTorch

torch

torch.arrange

torch.arange 作用类似于 Python 的 range()，用于生成一个连续数值的一维向量，start是起始值（默认 0），end是结束值（不包含端点），step是步长（默认为1）

tensor = torch.arange(start=0, end=10, step=1, dtype=None, device=None)
print(tensor) # tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

torch.expand

expand()方法用于扩展张量的维度，但不会复制数据，而是通过广播机制让多个维度共享相同的数据，使用如下：

sizes = x.shape
y = x.expand(*sizes) # （-1 代表保持原维度）

正余弦位置编码实现

对于位置 $pos$和维度索引 $i$，位置编码的计算公式为：

\begin{aligned}
PE_{(pos, 2i)} &= \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) \\
PE_{(pos, 2i+1)} &= \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)
\end{aligned}

Torch的代码实现如下所示：

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        
        # 生成位置编码矩阵 [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # [max_len, 1]
        
        # 计算div_term：10000^(2i/d_model)的倒数（对数空间计算避免溢出）
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-math.log(10000.0) / d_model
        ))  # [d_model/2]
        
        # 填充偶数和奇数维度
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数维度：sin(pos/10000^(2i/d))
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数维度：cos(pos/10000^(2i/d))
        
        # 注册为缓冲区（不参与训练）
        self.register_buffer('pe', pe.unsqueeze(0))  # [1, max_len, d_model]
        
    def forward(self, x):
        """
        x: 输入张量 [batch_size, seq_len, d_model]
        返回: [batch_size, seq_len, d_model]
        """
        return x + self.pe[:, :x.size(1)]  # 自动广播相加

正余弦编码的特点：固定不变（没有可学习参数）；低维编码变化快，高维编码变化慢，这是因为维度索引 $i$ 越大，则 $\frac{2i}{d_{\text{model}}}$ 越大，对应地三角函数的频率 $\omega$ 也就越小，因此三角函数的变化周期 $T=\frac{2 \pi}{\omega}$ 也就越慢