PyTorch Warmup

语法热身：换轴和合并维度怎么写

下面用一个“时间序列”小例子练语法，不使用 notebook 里的图像 shape。

# x: [batch, time, feature]
x = torch.randn(2, 4, 3)

# 把 feature 放到 time 前面: [batch, feature, time]
y = x.permute(0, 2, 1)

# 把 batch 和 time 合并: [batch * time, feature]
z = x.reshape(2 * 4, 3)

# einops 写法：名字就是维度含义
z2 = rearrange(x, "b t f -> (b t) f")

读法：permute 是“重新排列轴”，reshape 是“重新分组元素”，einops 是“把变形规则写成句子”。

举一反三：每次变形前写一行 shape 合同

# 输入合同: x [B,C,H,W]
# 目标合同: tokens [B,H*W,C]
# 数量守恒: C*H*W == (H*W)*C

tokens = x.permute(0, 2, 3, 1).reshape(B, H * W, C)
assert tokens.shape == (B, H * W, C)

shape 合同只能保证外形；最后仍要用一个小张量或 allclose 检查元素顺序。

语法热身：Embedding 和高级索引

这里用“颜色编号查颜色向量”的例子，不使用 notebook 的 vocab/hidden 设置。

# 4 种颜色，每种颜色用 3 个数字表示
palette = torch.tensor([
    [1.0, 0.0, 0.0],  # red
    [0.0, 1.0, 0.0],  # green
    [0.0, 0.0, 1.0],  # blue
    [1.0, 1.0, 0.0],  # yellow
])

color_ids = torch.tensor([[2, 0], [3, 1]])

# 高级索引：每个 id 替换成对应行
colors = palette[color_ids]  # shape: [2, 2, 3]

# nn.Embedding 的调用方式也一样：传入 id，返回向量
emb = nn.Embedding(num_embeddings=4, embedding_dim=3)
vectors = emb(color_ids)

读法：weight[input_ids] 会把 input_ids 的每个整数当成“行号”，输出 shape 会在 input_ids 后面追加 embedding_dim。

举一反三：用高级索引验证 nn.Embedding

official = emb_layer(input_ids)
manual = emb_layer.weight[input_ids]

assert official.shape == (*input_ids.shape, emb_layer.embedding_dim)
assert torch.allclose(official, manual)

这组对照把“Embedding 是查表”从比喻变成了可运行证据。

语法热身：自定义 autograd 函数长什么样

这里用“平方函数”的 backward 练结构，不直接写 notebook 的 Linear+ReLU。

class SquareFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x * x

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        grad_x = grad_output * 2 * x
        return grad_x

# 调用自定义 Function 要用 .apply
y = SquareFunction.apply(x)

读法：ctx.save_for_backward 像是在前向路上放书签；backward 里先取回书签，再按链式法则乘上传回来的 grad_output。

举一反三：用 autograd 当裁判

# custom 与 reference 使用同一份输入和参数
loss_custom.backward()
loss_reference.backward()

assert torch.allclose(custom_grad_x, reference_grad_x, atol=1e-5)
assert torch.allclose(custom_grad_w, reference_grad_w, atol=1e-5)

不要只比 forward 输出；自定义 Function 的核心验收对象是每一个输入梯度。