Amazon ML Coding：PyTorch Training Loop，更新与评估边界怎么写？

题目重点：train/eval、清梯度、前向loss反传更新和按样本加权指标，用手算SGD独立验证。

题意与假设

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

模型输出未softmax的logits[B,C]，target为long[B]，CPU默认，model须在目标device后再创建optimizer。
示例无AMP/累积/DDP/scheduler，调用方提供可迭代(features,targets) batches；空batch跳过、总样本0报错。
train_epoch返回各批更新前loss的样本加权均值，不是最终模型在全训练集重新计算的loss；evaluate返回固定模型loss/accuracy，并使模型保持eval模式。

解题路线

epoch开始model.train启用dropout/BN训练行为；每批move输入/target，zero_grad(set_to_none=True)，前向CE、backward、step，顺序必须解释。若把zero_grad放backward后step前，会清掉刚算梯度。
cross_entropy读logits而非概率。loss.detach().item用于日志避免把整个计算图留列表；参数更新前batch loss按batch大小累加，最后除总样本，末批小不等权。
evaluate同时model.eval与no_grad：前者改变dropout/BN行为，后者禁梯度记录，二者职责不同。全体正确数量/全体样本算accuracy，不平均各batch accuracy。
device迁移在optimizer构造前由调用方完成，dataset/label范围与模型输出类数对应。训练/验证按合同分离，不在eval里step或用validation标签反传。
加入梯度累积时先定义平均/归一化、何时step与清grad，最后不足累积批怎么处理；AMP需scaler/反缩放，scheduler跟真实更新。这些是后续扩展而非代码里隐藏支持。
验证不是只看loss下降：用NumPy softmax解析梯度手算Linear两次SGD更新，第二次检查没有旧grad累积；eval按不均匀末批与固定dropout模式核指标。

手动走一遍

五样本分4+1批时，batch均值loss分别1与3，正确epoch加权(4×1+1×3)/5=1.4，而不是(1+3)/2=2。随机Linear模型600次SGD步骤与解析更新、300组加权eval独立核对；dropout eval重复一致、空数据拒绝检查通过。

Python 参考实现

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

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

import torch
import torch.nn.functional as F


def train_epoch(model, batches, optimizer, device='cpu'):
    """Classification logits [B,C], long targets [B]; no AMP/accumulation/DDP."""
    model.train()
    total_loss, samples = 0.0, 0
    for features, targets in batches:
        features, targets = features.to(device), targets.to(device)
        if len(targets) == 0:
            continue
        optimizer.zero_grad(set_to_none=True)
        logits = model(features)
        loss = F.cross_entropy(logits, targets)
        loss.backward()
        optimizer.step()
        count = len(targets)
        total_loss += loss.detach().item() * count
        samples += count
    if samples == 0:
        raise ValueError('no nonempty training batches')
    return total_loss / samples


@torch.no_grad()
def evaluate(model, batches, device='cpu'):
    model.eval()
    total_loss, correct, samples = 0.0, 0, 0
    for features, targets in batches:
        features, targets = features.to(device), targets.to(device)
        if len(targets) == 0:
            continue
        logits = model(features)
        total_loss += F.cross_entropy(logits, targets, reduction='sum').item()
        correct += (logits.argmax(dim=-1) == targets).sum().item()
        samples += len(targets)
    if samples == 0:
        raise ValueError('no nonempty evaluation batches')
    return total_loss / samples, correct / samples

复杂度

每epoch成本为各批model前向/反传/optimizer更新，指标聚合O(样本数)，不会保留全epoch计算图；optimizer状态依参数量。eval不保存反传图，仍需模型前向算力。

易错点与追问

容易失分的边界

忘zero_grad导致跨批累积。
backward后清grad再step。
eval没设model.eval或no_grad。
把末批与满批等权。
日志保存loss tensor保留图。
softmax后再CE。

继续深挖

AMP溢出跳步如何处理scheduler？
最后不足梯度累积批怎样归一化？
恢复checkpoint须存哪些状态？
DDP怎样加权聚合各卡样本指标？