TikTok 神经网络训练与推理：尺寸、Batch、梯度和评估

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 实际尺寸和超参未披露；训练示例输入d、单ReLU隐藏h、C类输出logits，C≥2。
• 使用PyTorch、CPU张量与SGD；数据和标签先按模型dtype/device准备，无AMP/DDP。任务若回归需换输出和loss。

解题路线

1. 先确认 X[B,d]、标签[B]整数类ID，第一层W1[h,d]、第二层W2[C,h]，输出[B,C]，loss直接接logits而非先softmax。
2. 每batch依次zero_grad、forward、CE、backward、step；少zero_grad会累积历史梯度，少step不会训练。epoch平均按样本数加权，尾batch较小时不能简单平均batch均值。
3. 推理使用eval和no_grad；二者分别控制层行为与梯度记录。示例无dropout/BN，但完整合同为以后加入它们留正确切换。
4. 训练集/验证集先分开，预处理只拟合训练数据；追踪loss、指标、梯度和错误样本，保存模型及预处理版本。不把一个epoch下降当泛化成功。

手动走一遍

d=3、h=4、C=2，B=5时，first输出[5,4]、last输出[5,2]，CE得到标量。预测输出概率[5,2]和类别[5]；无有效batch明确拒绝，不能返回误导性0损失。

Python 参考实现

下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。

代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。

import torch
import torch.nn.functional as F


class MLP(torch.nn.Module):
    """One hidden ReLU layer, multiclass logits. Dimensions are explicit inputs."""
    def __init__(self, inputs, hidden, classes):
        super().__init__()
        if inputs < 1 or hidden < 1 or classes < 2:
            raise ValueError('positive dimensions and >=2 classes required')
        self.first = torch.nn.Linear(inputs, hidden)
        self.last = torch.nn.Linear(hidden, classes)
    def forward(self, x):
        return self.last(F.relu(self.first(x)))


def train_epoch(model, batches, optimizer):
    model.train()
    total, count = 0.0, 0
    for x, y in batches:
        if len(y) == 0:
            continue
        optimizer.zero_grad(set_to_none=True)
        logits = model(x)
        loss = F.cross_entropy(logits, y)
        loss.backward()
        optimizer.step()
        total += loss.detach().item() * len(y)
        count += len(y)
    if count == 0:
        raise ValueError('nonempty batches required')
    return total / count


@torch.no_grad()
def predict(model, x):
    model.eval()
    probabilities = model(x).softmax(dim=-1)
    return probabilities, probabilities.argmax(dim=-1)

复杂度

参数量dh+h+hC+C；单batch主乘法 O(Bdh+BhC)，激活 O(Bh+BC)，参数/梯度/优化器状态另外计。显存还含输入与框架临时量。

易错点与追问

容易失分的边界

• CE前额外softmax。
• 只写网络结构漏反向更新。
• eval替代no_grad。
• 训练验证数据泄漏。

继续深挖

• 多隐藏层怎样逐层推导shape？
• 如何加入早停、gradient accumulation和checkpoint？