Torch

pytorch

1. cross_entropy

输入 shape 约定：(B, C, *)，其中 dim=1 对应类别维 C。

给定 logits $z \in \mathbb{R}^{N \times C}$（$N$ 个样本，$C$ 个类别），标量 target $y \in \{0,1,\dots,C-1\}^N$，标准的 cross entropy loss 为：

$$ \mathcal{L} = -\frac{1}{N} \sum_{i=1}^{N} \log \frac{e^{\,z_{i, y_i}}}{\sum_{c=1}^{C} e^{\,z_{i, c}}} $$

底层等价于两步：

log_softmax(z, dim=1) — 沿类别维做 softmax 再取 log
nll_loss(log_softmax_output, target) — 取 target 对应位置的值，取负，求平均

常见例子：input 形状为 $(B, T, V)$（$V = \text{vocab\_size}$），target 为 $(B, T)$。计算 loss 时先做 .view：
$$ > \text{logits: } (B, T, V) \;\xrightarrow{\text{view}}\; (B \times T, V) > $$
$$ > \text{targets: } (B, T) \;\xrightarrow{\text{view}}\; (B \times T) > $$
target 的每一行都是一个整数 token id，即对应的 $y_i$，然后 F.cross_entropy(logits, targets) 直接算 loss。

数值稳定性

直接计算 $e^{z}$ 有两个问题：

上溢：$z$ 过大时 $e^z \to +\infty$
下溢：$z$ 过小时 $e^z \to 0$，随后 $\log(0) \to -\infty$

工程实现统一使用 LogSumExp trick：

$$ \log \sum_{j} e^{x_j} = m + \log \sum_{j} e^{\,x_j - m}, \quad m = \max_j x_j $$

即先减去每行最大值，再做 exp-sum-log。这样 $e^{x_j - m} \in (0, 1]$，exp 不会爆，log 也不会碰到 0。

对应到具体代码（log_softmax 内部已经做了这一步）：

$$ \text{log\_softmax}(z_i) = (z_i - \max z) - \log \sum_{c} e^{\,z_c - \max z} $$

Temperature（温度系数）用于调节分布锐度，与数值稳定性无关：
$$P(c) = \frac{e^{\,z_c / T}}{\sum_j e^{\,z_j / T}}$$
$T \to 0$ 趋近 one-hot（更确定），$T \to \infty$ 趋近均匀分布（更平滑）。

2. nn.Module

并不是所有 PyTorch 类都继承 nn.Module。判断标准：是否持有可训练参数（nn.Parameter）。

继承 `nn.Module`	不继承
`nn.Linear`, `nn.Embedding`, `nn.Conv2d`	`torch.Tensor`（张量本身）
`nn.LayerNorm`, `nn.BatchNorm`	`torch.optim.Adam`, `SGD`（优化器）
`nn.Transformer`, `nn.LSTM`	`DataLoader`, `Dataset`
自己定义的 model	`nn.functional.*`（函数式 API）
`nn.Sequential`, `nn.ModuleList`	`torch.cuda.amp.GradScaler`

规律：存参数、有 forward → nn.Module。优化器虽然更新参数但不持有参数（只持有引用），Tensor 是数据容器，DataLoader 管数据流，都不需要继承。

3. register_buffer

nn.Module 里注册一个不参与梯度更新但要跟着模型一起动的张量。

	`nn.Parameter`	`register_buffer`	普通 `self.x`
参与梯度	是	否	否
`model.to('cuda')` 自动迁移	是	是	否
存入 `state_dict`	是	是	否
优化器更新	是	否	否

典型用途：transformer 里的 causal mask、attention bias、RoPE 的 cos/sin 表——推理必须用、但不参与优化的东西。

self.register_buffer('bias', torch.tril(torch.ones(block_size, block_size))
                                 .view(1, 1, block_size, block_size))

4. nn.Parameter()

当我们自定义一个net或者组件的时候如果它是需要被优化器更新的,一定不要忘记注册为net的paramter
而像我们常写的类都是继承自nn.Module的,当然如果不要更新的话别忘了设置 required_grad = False

5. torch中设置时间的步骤

由于gpu和cpu是异步进行的, 所以

6. torch中的精度

torch中默认为fp32的精度, 一般来说我们用int8 来进行推理 tf32的精度对训练也没有什么影响, 所以我们可以通过