【python】高级切片语法/多维张量索引

整体上下文

# 核心代码
pe[0, :, 0::2] = torch.sin(position * div_term)    # even index
pe[0, :, 1::2] = torch.cos(position * div_term)    # odd index

语法详解

1. 多维张量索引语法

pe[0, :, 0::2]

这是一个三维张量索引，各部分含义：

• pe: 三维张量，形状 (batch_size, sequence_length, d_model)
• [0, :, 0::2]: 索引操作符，包含三个维度索引

维度详解：

• 第一个索引 0: 批次维度（batch dimension）
  • 获取第0个批次（如果批次大小为1，则索引0有效）
• 第二个索引 :: 序列长度维度（sequence dimension）
  • : 表示取该维度的所有元素（所有位置）
  • 相当于 0:max_len
• 第三个索引 0::2: 特征维度（feature dimension）
  • 步长切片（step slicing）

2. 步长切片语法 start:stop:step

# 基本语法
[start:stop:step]  # 从start开始，到stop结束（不包含），步长为step# 示例
pe[0, :, 0::2]  # 从索引0开始，步长为2（所有偶数索引）
pe[0, :, 1::2]  # 从索引1开始，步长为2（所有奇数索引）

详细说明：

• 0::2:

  • 0: 起始索引
  • 省略的: : 表示到该维度结束
  • 2: 步长（每2个取一个）
  • 相当于 0, 2, 4, 6, 8, ...

• 1::2:

  • 1: 起始索引
  • 2: 步长
  • 相当于 1, 3, 5, 7, 9, ...

3. 广播机制的应用

torch.sin(position * div_term)

这里涉及广播乘法：

• position: 形状 (max_len, 1)
• div_term: 形状 (d_model/2,)

广播过程：

1. position 形状变为 (max_len, 1)
2. div_term 形状变为 (1, d_model/2)
3. 结果形状：(max_len, d_model/2)

# 具体示例
max_len = 5
d_model = 8  # 所以 d_model/2 = 4position = torch.tensor([[0], [1], [2], [3], [4]])  # (5, 1)
div_term = torch.tensor([1.0, 0.1, 0.01, 0.001])    # (4,)# 广播后相乘
result = position * div_term  # 形状: (5, 4)
# 等价于：
# [[0*1.0, 0*0.1, 0*0.01, 0*0.001],
#  [1*1.0, 1*0.1, 1*0.01, 1*0.001],
#  [2*1.0, 2*0.1, 2*0.01, 2*0.001],
#  [3*1.0, 3*0.1, 3*0.01, 3*0.001],
#  [4*1.0, 4*0.1, 4*0.01, 4*0.001]]

4. 赋值操作的形状匹配

pe[0, :, 0::2] = torch.sin(position * div_term)

形状分析：

• 左侧 pe[0, :, 0::2]:

  • pe 形状: (1, max_len, d_model)
  • pe[0, :, 0::2] 形状: (max_len, d_model/2)
    • 第0维: max_len（所有序列位置）
    • 第1维: d_model/2（偶数索引的特征）

• 右侧 torch.sin(position * div_term):

  • 形状: (max_len, d_model/2)

形状完全匹配，可以赋值。