PyTorch小批次仅对部分样本反向传播无加速的原因是什么

• 核心开销已经在全量前向阶段完成：你当前的执行流程是先对完整小批次做前向传播、计算所有样本的损失，再做topk筛选。整个小批次的前向计算（尤其是DenseNet121的多层卷积运算）已经占用了90%以上的训练耗时，后续筛选和反向阶段的开销占比极低，自然看不到明显提速。
• 反向传播不会自动跳过未选中样本的计算：PyTorch的autograd是基于预先构建的完整计算图执行反向传播的，你调用topk只是筛选了最终的损失值，没有改变前面已经构建好的全量计算图结构。未被选中的样本对应的损失梯度为0，但反向传播依然会遍历整个计算图的所有节点完成梯度计算，不会自动跳过部分节点，因此反向阶段也不会有明显的速度提升。
• topk操作本身不具备关闭部分样本autograd的功能：topk是可微的张量索引操作，仅会对选中的数值传递正常梯度、未选中的数值传递0梯度，不会主动截断计算图或者关闭部分节点的梯度跟踪，因此不会减少反向传播的计算量。

• 方案1：无梯度预筛选+小样本前向反向
  调整执行流程，先在无梯度环境下完成全批次损失计算和样本筛选，再仅对选中的样本做正式的前向和反向传播，即可减少前后向的计算量：

# 第一步：无梯度预计算所有样本损失，筛选topk下标
with torch.no_grad():
    model.eval() # 避免预计算影响BN等层的运行统计量
    outputs = model(inputs)
    pre_loss = loss_fn(outputs, targets, reduction='none')
    r = max(round(0.6 * inputs.shape[0]), 1)
    _, topk_idx = torch.topk(pre_loss, r, largest=False, sorted=False, dim=0)
# 第二步：仅对选中样本做正式前向反向
model.train()
selected_inputs = inputs[topk_idx]
selected_targets = targets[topk_idx]
selected_outputs = model(selected_inputs)
loss = loss_fn(selected_outputs, selected_targets).mean()
loss.backward()
optimizer.step()

注意如果你的模型里有BN、Dropout等训练/评估状态不同的层，预计算阶段要把模型切到eval模式，避免预计算过程修改层的运行统计量。

• 方案2：如果不需要提速仅要训练效果，可保留原有逻辑
  如果你调整损失的目标只是重点学习低损失样本，不需要提速，那么你原来的代码逻辑是正确的，不需要额外修改，只是确实不会带来训练速度提升。

内容的提问来源于stack exchange，提问作者user2454869